diff --git a/RELEASE.md b/RELEASE.md index 38afa6c8b11cbe..ff95217c3bb665 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -14,7 +14,8 @@ ### Major Features and Improvements -* +* `tf.lite` + * Adds blockwise quantization support in the converter and `quant_spec` metadata in `FullyConnectedOptions`. ### Bug Fixes and Other Changes diff --git a/tensorflow/compiler/mlir/lite/BUILD b/tensorflow/compiler/mlir/lite/BUILD index c5da7b7e7139e5..cdb76a3f91a6aa 100644 --- a/tensorflow/compiler/mlir/lite/BUILD +++ b/tensorflow/compiler/mlir/lite/BUILD @@ -1155,6 +1155,7 @@ cc_library( "@llvm-project//mlir:IR", "@llvm-project//mlir:Pass", "@llvm-project//mlir:Support", + "@llvm-project//mlir:TransformUtils", ], ) @@ -1552,6 +1553,7 @@ cc_library( srcs = [ "transforms/decompose_hybrid_quantization.cc", "transforms/default_quant_params.cc", + "transforms/fold_stablehlo_constant_transforms_pass.cc", "transforms/generated_post_quantize.inc", "transforms/generated_quantize.inc", "transforms/lower_quant_annotations_helper.cc", diff --git a/tensorflow/compiler/mlir/lite/converter_gen.cc b/tensorflow/compiler/mlir/lite/converter_gen.cc index b08972c0301f33..df8c3c5d111e06 100644 --- a/tensorflow/compiler/mlir/lite/converter_gen.cc +++ b/tensorflow/compiler/mlir/lite/converter_gen.cc @@ -62,11 +62,11 @@ llvm::cl::opt action( "Generate TFLite runtime verifiers"))); // Returns the associated option name for the given op definition. -static inline std::string GetOperatorOptionName(const Record &def) { +static inline std::string GetOperatorOptionName(const Record& def) { assert(def.getName().starts_with("TFL_") && "unexpected op prefix"); assert(def.getName().ends_with("Op") && "unexpected op suffix"); - auto *custom_option = dyn_cast(def.getValueInit("customOption")); + auto* custom_option = dyn_cast(def.getValueInit("customOption")); std::ostringstream oss; if (custom_option) oss << custom_option->getValue().str(); @@ -90,7 +90,7 @@ static inline bool IsLstmOp(const StringRef op_name) { return op_name.take_back(6) == "LSTMOp"; } -static int HasOptions(const Record &def) { +static int HasOptions(const Record& def) { if (def.getValueAsBit("hasOptions")) { return 1; } @@ -100,13 +100,13 @@ static int HasOptions(const Record &def) { return 0; } -static void EmitOptionBuilders(const RecordKeeper &record_keeper, - const std::vector &defs, - raw_ostream *ostream) { - raw_ostream &os = *ostream; +static void EmitOptionBuilders(const RecordKeeper& record_keeper, + const std::vector& defs, + raw_ostream* ostream) { + raw_ostream& os = *ostream; const auto attr_type = record_keeper.getClass("Attr"); - for (const auto *def : defs) { + for (const auto* def : defs) { const int has_options = HasOptions(*def); // TFLite ops without options are skipped over. if (!has_options) { @@ -125,11 +125,11 @@ static void EmitOptionBuilders(const RecordKeeper &record_keeper, // Construct all the builder option needed. SmallVector options; // Add options due to attributes (not-derived). - auto *arg_values = def->getValueAsDag("arguments"); + auto* arg_values = def->getValueAsDag("arguments"); mlir::tblgen::Operator op(*def); for (unsigned i = 0, e = arg_values->getNumArgs(); i != e; ++i) { auto arg = arg_values->getArg(i); - const auto *arg_def = dyn_cast(arg); + const auto* arg_def = dyn_cast(arg); if (!arg_def) continue; if (arg_def->getDef()->isSubClassOf(attr_type)) { // This binds the name of the attribute in the TD file with the name @@ -156,8 +156,8 @@ static void EmitOptionBuilders(const RecordKeeper &record_keeper, } // Add options due to derived attributes. - for (const auto &val : def->getValues()) { - if (auto *record = dyn_cast(val.getType())) { + for (const auto& val : def->getValues()) { + if (auto* record = dyn_cast(val.getType())) { if (record->isSubClassOf(attr_type)) { if (record->getClasses().size() != 1) { PrintFatalError( @@ -174,7 +174,7 @@ static void EmitOptionBuilders(const RecordKeeper &record_keeper, } os << " tflite::" << tflite_option_name << "Builder b(*fbb);\n"; - for (const auto &option : options) + for (const auto& option : options) os << formatv(" b.add_{0}(std::move({0}));\n", option); os << " return b.Finish();\n}\n"; } @@ -187,11 +187,11 @@ static void EmitOptionBuilders(const RecordKeeper &record_keeper, // arguments that depend on op definitions should be auto-generated and then // operator should be built by the caller because it does not require // auto-generation. -static void EmitOperatorBuilders(const std::vector &defs, - raw_ostream *ostream) { - raw_ostream &os = *ostream; +static void EmitOperatorBuilders(const std::vector& defs, + raw_ostream* ostream) { + raw_ostream& os = *ostream; - for (const auto *def : defs) { + for (const auto* def : defs) { StringRef op_name = def->getName().drop_front(4); const bool has_intermediates = op_name.take_back(6) == "LSTMOp"; @@ -257,7 +257,7 @@ static void EmitOperatorBuilders(const std::vector &defs, } } -static inline std::string GetOperatorName(const Record &def) { +static inline std::string GetOperatorName(const Record& def) { auto name = def.getValueAsString("opName"); // Special case for basic_lstm. if (name == "basic_lstm") { @@ -275,14 +275,14 @@ static inline std::string GetOperatorName(const Record &def) { // // TODO(hinsu): Consider converting this to a static constant associative // container instead of a series of if conditions, if required. -static void EmitGetBuiltinOpCode(const std::vector &defs, - raw_ostream *ostream) { - raw_ostream &os = *ostream; +static void EmitGetBuiltinOpCode(const std::vector& defs, + raw_ostream* ostream) { + raw_ostream& os = *ostream; os << "std::optional " "mlir::GetBuiltinOpCode(mlir::Operation* op) {\n"; - for (const auto *def : defs) { + for (const auto* def : defs) { StringRef op_name = def->getName().drop_front(4); auto operator_name = GetOperatorName(*def); os << " if (isa(op))\n" @@ -304,20 +304,20 @@ static void EmitGetBuiltinOpCode(const std::vector &defs, // ... // return {0, 0}; // } -static void EmitOperandNumbers(const RecordKeeper &record_keeper, - const std::vector &defs, - raw_ostream *ostream) { - raw_ostream &os = *ostream; +static void EmitOperandNumbers(const RecordKeeper& record_keeper, + const std::vector& defs, + raw_ostream* ostream) { + raw_ostream& os = *ostream; const auto attr_type = record_keeper.getClass("Attr"); const auto optional_tensor = record_keeper.getClass("TFL_TensorOfOrNone"); os << "llvm::MinMax mlir::OperandNumbersMinMax(llvm::StringRef op_name) {\n"; - for (const auto *def : defs) { + for (const auto* def : defs) { auto op_name = def->getValueAsString("opName"); int tail_optional_tensor = 0, tensor_number_max = 0; - auto *arg_values = def->getValueAsDag("arguments"); + auto* arg_values = def->getValueAsDag("arguments"); for (int i = 0, e = arg_values->getNumArgs(); i < e; ++i) { auto arg = arg_values->getArg(i); - auto *arg_def = dyn_cast(arg); + auto* arg_def = dyn_cast(arg); if (!arg_def) continue; if (!arg_def->getDef()->isSubClassOf(attr_type)) { tensor_number_max++; @@ -351,9 +351,9 @@ static void EmitOperandNumbers(const RecordKeeper &record_keeper, // const std::vector& intermediates, // flatbuffers::FlatBufferBuilder *fbb, // std::optional debug_metadata_index); -static void EmitBuildOperator(const std::vector &defs, - raw_ostream *ostream) { - raw_ostream &os = *ostream; +static void EmitBuildOperator(const std::vector& defs, + raw_ostream* ostream) { + raw_ostream& os = *ostream; // Signature os << "std::optional>\n" @@ -365,7 +365,7 @@ static void EmitBuildOperator(const std::vector &defs, "flatbuffers::FlatBufferBuilder *fbb," "std::optional debug_metadata_index) {\n"; - for (const auto *def : defs) { + for (const auto* def : defs) { StringRef op_name = def->getName().drop_front(4); // Try to cast to each op case and call the corresponding op builder @@ -391,9 +391,9 @@ static void EmitBuildOperator(const std::vector &defs, // where id is an empty string if builtin_options_id is 1, or builtin_options_id // otherwise. static void EmitBuiltinOptionsToAttributes( - const RecordKeeper &record_keeper, const std::vector &defs, - raw_ostream *ostream, const int builtin_options_id) { - raw_ostream &os = *ostream; + const RecordKeeper& record_keeper, const std::vector& defs, + raw_ostream* ostream, const int builtin_options_id) { + raw_ostream& os = *ostream; const std::string builtin_options_suffix = [&] { switch (builtin_options_id) { @@ -415,7 +415,7 @@ static void EmitBuiltinOptionsToAttributes( "llvm::SmallVectorImpl &attributes) {\n"; const auto attr_type = record_keeper.getClass("Attr"); - for (const auto *def : defs) { + for (const auto* def : defs) { const int has_options = HasOptions(*def); if (has_options != builtin_options_id) { continue; @@ -429,10 +429,10 @@ static void EmitBuiltinOptionsToAttributes( os << formatv(" if(const auto *op = op_union.As{0}()) {{\n", option_name); // We only care about options that are in arguments - auto *arg_values = def->getValueAsDag("arguments"); + auto* arg_values = def->getValueAsDag("arguments"); for (unsigned i = 0, e = arg_values->getNumArgs(); i != e; ++i) { auto arg = arg_values->getArg(i); - const auto *arg_def = dyn_cast(arg); + const auto* arg_def = dyn_cast(arg); if (!arg_def) continue; if (arg_def->getDef()->isSubClassOf(attr_type)) { StringRef arg_name = arg_values->getArgNameStr(i); @@ -449,6 +449,21 @@ static void EmitBuiltinOptionsToAttributes( } } + for (const auto& val : def->getValues()) { + if (auto* record = dyn_cast(val.getType())) { + if (record->isSubClassOf(attr_type) && + record->getClasses().size() == 1 && + record->getClasses()[0]->getName() == "DerivedQuantSpecAttr") { + os << formatv( + " if (auto attr = BuildDerivedQuantSpecAttr(op->{0}, " + "builder))\n" + " attributes.emplace_back(builder.getNamedAttr(" + "mlir::TFL::kQuantSpecAttrName, attr));\n", + val.getName()); + } + } + } + os << " return;\n"; os << " }\n"; } @@ -463,14 +478,14 @@ static void EmitBuiltinOptionsToAttributes( // The function below has a non-constant reference as that is required by LLVM's // TableGenMain. // NOLINTNEXTLINE -static bool OperatorWritersMain(raw_ostream &os, const RecordKeeper &records) { +static bool OperatorWritersMain(raw_ostream& os, const RecordKeeper& records) { emitSourceFileHeader("MLIR TFLite FlatBuffer Builders", os); // Retrieve all the definitions derived from TFL_Op and sort by record name. - std::vector defs = records.getAllDerivedDefinitions("TFL_Op"); + std::vector defs = records.getAllDerivedDefinitions("TFL_Op"); llvm::sort(defs, LessRecord()); - for (const auto *def : defs) { + for (const auto* def : defs) { // TFLite ops in the .td file are expected to follow the naming convention: // TFL_Op. // The generated TFLite op C++ class should be TFL::Op. @@ -501,15 +516,15 @@ static bool OperatorWritersMain(raw_ostream &os, const RecordKeeper &records) { return false; } -static void GenOperandResultVerifier(raw_ostream &os, - llvm::ArrayRef values, +static void GenOperandResultVerifier(raw_ostream& os, + llvm::ArrayRef values, StringRef valueKind) { mlir::tblgen::FmtContext fctx; bool first = true; - for (const auto &static_value : llvm::enumerate(values)) { - auto *definit = llvm::cast(static_value.value()); - auto *val = definit->getDef()->getValue("tflRuntimeTypePredicate"); + for (const auto& static_value : llvm::enumerate(values)) { + auto* definit = llvm::cast(static_value.value()); + auto* val = definit->getDef()->getValue("tflRuntimeTypePredicate"); if (!val) continue; // Create code block on first type to verify. @@ -550,16 +565,16 @@ static void GenOperandResultVerifier(raw_ostream &os, } // NOLINTNEXTLINE -static bool RuntimeVerifierWriterMain(raw_ostream &os, - const RecordKeeper &records) { +static bool RuntimeVerifierWriterMain(raw_ostream& os, + const RecordKeeper& records) { emitSourceFileHeader("MLIR TFLite Runtime Verifiers", os); // Retrieve all the definitions derived from TFL_Op and sort by record name. - std::vector defs = records.getAllDerivedDefinitions("Op"); + std::vector defs = records.getAllDerivedDefinitions("Op"); llvm::sort(defs, LessRecord()); // Iterate through all the ops defined. - for (const auto *def : defs) { + for (const auto* def : defs) { mlir::tblgen::Operator op(*def); if (!op.getTrait("TflRuntimeVerifyOpInterface::Trait")) continue; @@ -571,7 +586,7 @@ static bool RuntimeVerifierWriterMain(raw_ostream &os, verify_ctx.addSubst("_op", "(*op)"); for (int i = 0, e = op.getNumOperands(); i < e; ++i) { - auto &value = op.getOperand(i); + auto& value = op.getOperand(i); // Skip from first variadic operands for now. Else getOperand index used // below doesn't match. if (value.isVariableLength()) break; @@ -579,7 +594,7 @@ static bool RuntimeVerifierWriterMain(raw_ostream &os, verify_ctx.addSubst(value.name, formatv("op->getOperand({0})", i)); } for (int i = 0, e = op.getNumResults(); i < e; ++i) { - auto &value = op.getResult(i); + auto& value = op.getResult(i); // Skip from first variadic results for now. Else getResult index used // below doesn't match. if (value.isVariableLength()) break; @@ -591,7 +606,7 @@ static bool RuntimeVerifierWriterMain(raw_ostream &os, GenOperandResultVerifier(os, def->getValueAsDag("results")->getArgs(), "result"); - for (auto &trait : op.getTraits()) { + for (auto& trait : op.getTraits()) { if (!trait.getDef().isSubClassOf("GenInternalOpTrait")) { continue; } @@ -600,7 +615,7 @@ static bool RuntimeVerifierWriterMain(raw_ostream &os, continue; } - auto *val = trait.getDef().getValue("tflRuntimePredicate"); + auto* val = trait.getDef().getValue("tflRuntimePredicate"); if (!val) continue; auto desc = trait.getDef().getValueAsString("tflRuntimeDescription"); @@ -643,7 +658,7 @@ static bool RuntimeVerifierWriterMain(raw_ostream &os, return false; } -int main(int argc, char **argv) { +int main(int argc, char** argv) { llvm::InitLLVM y(argc, argv); llvm::cl::ParseCommandLineOptions(argc, argv); if (action == ActionType::OpConv) diff --git a/tensorflow/compiler/mlir/lite/core/api/flatbuffer_conversions.cc b/tensorflow/compiler/mlir/lite/core/api/flatbuffer_conversions.cc index 2c4b62b39ec179..7747bc5368564b 100644 --- a/tensorflow/compiler/mlir/lite/core/api/flatbuffer_conversions.cc +++ b/tensorflow/compiler/mlir/lite/core/api/flatbuffer_conversions.cc @@ -1871,6 +1871,11 @@ absl::Status ParseFullyConnected(const Operator* op, schema_params->asymmetric_quantize_inputs(); TFL_FILE_ENSURE_STATUS(ConvertTensorType( schema_params->quantized_bias_type(), ¶ms->quantized_bias_type)); + if (const auto* quant_spec = schema_params->quant_spec()) { + // Borrowed from the model buffer; see TfLiteFullyConnectedParams. + params->quant_spec = quant_spec->data(); + params->quant_spec_size = static_cast(quant_spec->size()); + } switch (schema_params->weights_format()) { case FullyConnectedOptionsWeightsFormat_DEFAULT: params->weights_format = kTfLiteFullyConnectedWeightsFormatDefault; diff --git a/tensorflow/compiler/mlir/lite/core/c/builtin_op_data.h b/tensorflow/compiler/mlir/lite/core/c/builtin_op_data.h index c580bf03cd3f59..c69daeddc44c2f 100644 --- a/tensorflow/compiler/mlir/lite/core/c/builtin_op_data.h +++ b/tensorflow/compiler/mlir/lite/core/c/builtin_op_data.h @@ -208,6 +208,18 @@ typedef struct { // Parameters for FullyConnected version 10 or above. // Used to determine the default value for the quantized bias. TfLiteType quantized_bias_type; + + // Parameters for FullyConnected version 15 or above. + // + // Opaque, flexbuffers-encoded description of a non-standard quantization + // contract; see `FullyConnectedOptions.quant_spec` in schema.fbs. Null when + // the op uses standard FullyConnected quantization semantics. + // + // This points into the model's buffer and is only valid for as long as the + // model is, in the same way as `TfLiteNode::custom_initial_data`. It is not + // owned by this struct. + const uint8_t* quant_spec; + int quant_spec_size; } TfLiteFullyConnectedParams; typedef enum { diff --git a/tensorflow/compiler/mlir/lite/flatbuffer_operator.cc b/tensorflow/compiler/mlir/lite/flatbuffer_operator.cc index 1262ffa2191be2..5394b3566a7cce 100644 --- a/tensorflow/compiler/mlir/lite/flatbuffer_operator.cc +++ b/tensorflow/compiler/mlir/lite/flatbuffer_operator.cc @@ -27,7 +27,6 @@ limitations under the License. #include "absl/strings/str_cat.h" #include "flatbuffers/buffer.h" // from @flatbuffers #include "flatbuffers/flatbuffer_builder.h" // from @flatbuffers -#include "flatbuffers/flatbuffers.h" // from @flatbuffers #include "flatbuffers/flexbuffers.h" // from @flatbuffers #include "flatbuffers/string.h" // from @flatbuffers #include "flatbuffers/vector.h" // from @flatbuffers @@ -52,11 +51,8 @@ limitations under the License. #include "tensorflow/compiler/mlir/lite/schema/mutable/schema_generated.h" #include "tensorflow/compiler/mlir/lite/schema/schema_utils.h" #include "tensorflow/compiler/mlir/lite/utils/convert_type.h" -#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h" #include "tensorflow/compiler/mlir/tensorflow/utils/dynamic_shape_utils.h" -#include "tensorflow/core/platform/errors.h" #include "tensorflow/core/platform/status.h" -#include "tsl/platform/status.h" namespace { @@ -252,6 +248,78 @@ static tflite::LSTMKernelType ConvertTFL_LSTMKernelTypeAttrForOptionWriter( llvm_unreachable("invalid lstm_kernel_type in conversion."); } +// Serializes one `quant_spec` value. Only the kinds that the specs actually +// use are supported; anything else is a bug in the pass that built the +// dictionary, and dropping it would silently change the op's meaning. +// +// The key is taken by `const std::string&` rather than `llvm::StringRef` +// because the `flexbuffers::Builder` key overloads call `strlen` on it, and a +// `StringRef` is not guaranteed to be null-terminated. +static void SerializeQuantSpecValue(const std::string& key, + mlir::Attribute value, + flexbuffers::Builder* fbb) { + if (auto str = mlir::dyn_cast(value)) { + fbb->String(key.c_str(), str.getValue().str()); + } else if (auto boolean = mlir::dyn_cast(value)) { + fbb->Bool(key.c_str(), boolean.getValue()); + } else if (auto integer = mlir::dyn_cast(value)) { + fbb->Int(key.c_str(), integer.getInt()); + } else if (auto fp = mlir::dyn_cast(value)) { + fbb->Double(key.c_str(), fp.getValueAsDouble()); + } else { + llvm::report_fatal_error(llvm::Twine("quant_spec key '") + key + + "' has an unsupported attribute kind"); + } +} + +static flatbuffers::Offset> +ConvertDerivedQuantSpecAttrForOptionWriter( + mlir::DictionaryAttr attr, flatbuffers::FlatBufferBuilder* builder) { + if (!attr || attr.empty()) return 0; + + flexbuffers::Builder fbb; + const size_t map_start = fbb.StartMap(); + for (const mlir::NamedAttribute& named : attr) { + SerializeQuantSpecValue(named.getName().str(), named.getValue(), &fbb); + } + fbb.EndMap(map_start); + fbb.Finish(); + return builder->CreateVector(fbb.GetBuffer()); +} + +static mlir::DictionaryAttr BuildDerivedQuantSpecAttr( + const std::vector& quant_spec, mlir::Builder builder) { + if (quant_spec.empty()) return nullptr; + + flexbuffers::Reference root = flexbuffers::GetRoot(quant_spec); + if (!root.IsMap()) return nullptr; + flexbuffers::Map map = root.AsMap(); + flexbuffers::TypedVector keys = map.Keys(); + flexbuffers::Vector values = map.Values(); + + llvm::SmallVector entries; + entries.reserve(keys.size()); + for (size_t i = 0; i < keys.size(); ++i) { + llvm::StringRef key = keys[i].AsKey(); + flexbuffers::Reference val = values[i]; + mlir::Attribute attr; + if (val.IsString()) { + attr = builder.getStringAttr(val.AsString().str()); + } else if (val.IsBool()) { + attr = builder.getBoolAttr(val.AsBool()); + } else if (val.IsIntOrUint()) { + attr = builder.getI64IntegerAttr(val.AsInt64()); + } else if (val.IsFloat()) { + attr = builder.getF32FloatAttr(val.AsFloat()); + } else { + llvm::report_fatal_error(llvm::Twine("quant_spec key '") + key + + "' has an unsupported flexbuffer type"); + } + entries.push_back(builder.getNamedAttr(key, attr)); + } + return builder.getDictionaryAttr(entries); +} + static mlir::Attribute BuildBoolAttr(bool value, mlir::Builder builder) { return builder.getBoolAttr(value); } diff --git a/tensorflow/compiler/mlir/lite/ir/tfl_ops.cc b/tensorflow/compiler/mlir/lite/ir/tfl_ops.cc index b8527afd3648ba..f286d667cdadc4 100644 --- a/tensorflow/compiler/mlir/lite/ir/tfl_ops.cc +++ b/tensorflow/compiler/mlir/lite/ir/tfl_ops.cc @@ -1971,6 +1971,121 @@ LogicalResult BatchMatMulOp::verify() { return success(); } +//===----------------------------------------------------------------------===// +// BlockwiseQuantizeOp / BlockwiseDequantizeOp +//===----------------------------------------------------------------------===// + +namespace { + +// Returns the shape of the scale/zero-point grid implied by tiling `type` with +// `block_shape`, or failure (after emitting a diagnostic on `op`) if the two +// are not compatible. +FailureOr> GetBlockGridShape(Operation* op, + ShapedType type, + ArrayAttr block_shape) { + if (!type.hasStaticShape()) { + return op->emitOpError("expects a statically shaped tensor, got ") << type; + } + ArrayRef shape = type.getShape(); + if (block_shape.size() != shape.size()) { + return op->emitOpError("expects block_shape of rank ") + << shape.size() << ", got " << block_shape; + } + + SmallVector grid_shape; + grid_shape.reserve(shape.size()); + for (auto [dim, dim_size] : llvm::enumerate(shape)) { + auto block_attr = mlir::dyn_cast(block_shape[dim]); + if (!block_attr) { + return op->emitOpError("expects an integer block_shape, got ") + << block_shape; + } + const int64_t block_size = block_attr.getInt(); + if (block_size <= 0) { + return op->emitOpError("expects a positive block_shape, got ") + << block_shape; + } + if (dim_size % block_size != 0) { + return op->emitOpError("expects dimension ") + << dim << " (" << dim_size << ") to be divisible by block_shape[" + << dim << "] (" << block_size << ")"; + } + grid_shape.push_back(dim_size / block_size); + } + return grid_shape; +} + +// Verifies that `value`, if present, is shaped like the block grid. A +// dimension of size 1 is accepted anywhere as a broadcast. +LogicalResult VerifyBlockGridOperand(Operation* op, Value value, + ArrayRef grid_shape, + StringRef name) { + if (!value || mlir::isa(value.getType())) return success(); + auto type = mlir::dyn_cast(value.getType()); + if (!type) return success(); + + if (type.getRank() != static_cast(grid_shape.size())) { + return op->emitOpError("expects ") + << name << " of rank " << grid_shape.size() << ", got " << type; + } + for (auto [dim, dim_size] : llvm::enumerate(type.getShape())) { + if (dim_size != 1 && dim_size != grid_shape[dim]) { + return op->emitOpError("expects ") + << name << " dimension " << dim << " to be 1 or " + << grid_shape[dim] << ", got " << type; + } + } + return success(); +} + +} // namespace + +LogicalResult BlockwiseQuantizeOp::verify() { + auto grid_shape = GetBlockGridShape( + *this, mlir::cast(getInput().getType()), getBlockShape()); + if (failed(grid_shape)) return failure(); + + if (failed(VerifyBlockGridOperand(*this, getScale(), *grid_shape, "scale")) || + failed(VerifyBlockGridOperand(*this, getZeroPoint(), *grid_shape, + "zero_point"))) { + return failure(); + } + + if (getScaleType() != + mlir::cast(getScale().getType()).getElementType()) { + return emitOpError("expects the scale element type to match scale_type (") + << getScaleType() << ")"; + } + return success(); +} + +LogicalResult BlockwiseDequantizeOp::verify() { + auto grid_shape = GetBlockGridShape( + *this, mlir::cast(getInput().getType()), getBlockShape()); + if (failed(grid_shape)) return failure(); + + if (failed( + VerifyBlockGridOperand(*this, getScales(), *grid_shape, "scales")) || + failed(VerifyBlockGridOperand(*this, getZeroPoints(), *grid_shape, + "zero_points"))) { + return failure(); + } + + if (getSymmetric() && getZeroPoints() && + !mlir::isa(getZeroPoints().getType())) { + auto zero_points = mlir::dyn_cast_or_null( + getZeroPoints().getDefiningOp() + ? getZeroPoints().getDefiningOp()->getAttrOfType( + "value") + : nullptr); + if (zero_points && !zero_points.isSplat()) { + return emitOpError( + "expects a per-tensor zero_point when symmetric is set"); + } + } + return success(); +} + //===----------------------------------------------------------------------===// // FullyConnectedOp //===----------------------------------------------------------------------===// diff --git a/tensorflow/compiler/mlir/lite/ir/tfl_ops.h b/tensorflow/compiler/mlir/lite/ir/tfl_ops.h index 89d1a5ed960275..6e51d2334aba99 100644 --- a/tensorflow/compiler/mlir/lite/ir/tfl_ops.h +++ b/tensorflow/compiler/mlir/lite/ir/tfl_ops.h @@ -55,6 +55,12 @@ class ControlType : public Type::TypeBase { static constexpr StringLiteral name = "tfl.control"; }; +// Discardable attribute naming a non-standard quantization contract that the +// operand types cannot express on their own. Read by the `quant_spec` derived +// attribute of `tfl.fully_connected`; see `FullyConnectedOptions.quant_spec` in +// schema.fbs. +inline constexpr StringLiteral kQuantSpecAttrName = "tfl.quant_spec"; + #include "tensorflow/compiler/mlir/lite/ir/tfl_ops_interface.h.inc" } // end namespace TFL diff --git a/tensorflow/compiler/mlir/lite/ir/tfl_ops.td b/tensorflow/compiler/mlir/lite/ir/tfl_ops.td index b281b8bee7748d..7a2778dc978f2a 100644 --- a/tensorflow/compiler/mlir/lite/ir/tfl_ops.td +++ b/tensorflow/compiler/mlir/lite/ir/tfl_ops.td @@ -58,6 +58,13 @@ def TFL_Control: Type($_self)">, "control">, // A type attribute containing the TensorType. def TensorTypeAttr : TypeAttrBase<"TensorType", "Tensor type attribute">; +// The element type of the scale that a blockwise quantize op computes at +// runtime. Only `f32` and `f8E8M0FNU` are meaningful today; the latter +// constrains the scale to a power of two. +def TFL_ScaleTypeAttr : TypeAttrBase<"mlir::Type", "scale element type attribute"> { + let constBuilderCall = "($0 ? ::mlir::TypeAttr::get($0) : ::mlir::TypeAttr())"; +} + //===----------------------------------------------------------------------===// // Derived shape attribute class. //===----------------------------------------------------------------------===// @@ -65,6 +72,11 @@ class DerivedShapeAttr : DerivedAttr<"ArrayRef", body>; class DerivedTFLiteTypeAttr : DerivedAttr<"tflite::TensorType", body, convert>; +// Reads the opaque quantization spec of an op; see +// `FullyConnectedOptions.quant_spec` in schema.fbs. Null when the op uses +// standard quantization semantics. +class DerivedQuantSpecAttr : DerivedAttr<"::mlir::DictionaryAttr", body>; + // TFL Runtime op trait predicate. class TFL_RuntimePredOpTrait : GenInternalOpTrait<"TFLRuntimeOpTrait"> { @@ -1140,6 +1152,21 @@ def TFL_FullyConnectedOp : TFL_Op<"fully_connected", [ OptionalAttr:$asymmetric_quantize_inputs ); + // Names a non-standard quantization contract that the operand types cannot + // express on their own; see `FullyConnectedOptions.quant_spec` in schema.fbs. + // A runtime that does not recognize the named spec must reject the op rather + // than fall back to the standard behavior. + // + // This is a derived attribute reading the discardable `tfl.quant_spec` + // rather than an entry in `arguments`, because adding an argument would + // change the generated builders and break every pass, in this tree and out + // of it, that constructs a fully_connected positionally. + DerivedQuantSpecAttr quant_spec = DerivedQuantSpecAttr<[{ + return (*this)->getAttrOfType<::mlir::DictionaryAttr>( + TFL::kQuantSpecAttrName); + }]>; + + // Depending on the weights format, this op can have one or two outputs. let results = (outs TFL_VariadicTensorOf<[F32, F16, QI8, QUI8, QI16, QUI16]>:$output @@ -4352,6 +4379,80 @@ def TFL_DequantizeOp: TFL_Op<"dequantize", [NoMemoryEffect]> { ]; } +// The two ops below model blockwise (sub-channel) quantization, which +// `!quant.uniform` cannot express. They are deliberately *not* part of the +// TFLite schema: they only exist between `LowerQuantAnnotations` and +// `FuseA4W2DynamicRangeFullyConnected`, which collapses the recognized +// patterns back into a single `tfl.fully_connected`. An op that survives to +// the exporter is a bug, and the exporter reports it as an unsupported op. +// +// Block layout follows the convention used throughout this dialect: the +// tensor is tiled by `block_shape`, one element of the scale/zero-point grid +// per tile, with the grid laid out in the same dimension order as the tensor. +// A grid dimension of size 1 broadcasts across that axis. + +// These two ops exist only inside the MLIR pipeline: they are declared with the +// raw `Op` base rather than `TFL_Op` so that `converter_gen` +// does not try to map them onto a `BuiltinOperator` that the schema does not +// have. `FuseA4W2DynamicRangeFullyConnected` collapses them into a single +// `tfl.fully_connected`; any that survive to export fail there. +def TFL_BlockwiseQuantizeOp : Op { + let summary = "Blockwise dynamic quantize operator"; + + let description = [{ + Computes a per-block scale (and, when asymmetric, zero point) from the + input at runtime and quantizes the input with them. The computed + parameters are returned alongside the quantized tensor so that a matching + `tfl.blockwise_dequantize` can undo the transformation. + + The scale of a block is `max(abs(block)) / (q_max + range_dilation)`, + where `q_max` is the largest representable value of the output storage + type. `range_dilation` widens the denominator so that outliers clip + instead of stretching the whole block; it defaults to 0. When + `scale_type` is `f8E8M0FNU`, the scale is rounded up to the next power of + two. + }]; + + let arguments = (ins + TFL_FpTensor:$input, + I64ArrayAttr:$block_shape, + TFL_ScaleTypeAttr:$scale_type, + BoolAttr:$symmetric, + DefaultValuedAttr:$range_dilation + ); + + let results = (outs + TensorOf<[TFL_I2, TFL_I4, I8, I16]>:$output, + TFL_TensorOf<[F8E8M0FNU, F16, BF16, F32]>:$scale, + TFL_TensorOfOrNone<[TFL_I2, TFL_I4, I8, I16, F32]>:$zero_point + ); + + let hasVerifier = 1; +} + +def TFL_BlockwiseDequantizeOp + : Op { + let summary = "Blockwise dequantize operator"; + + let description = [{ + Reconstructs floating point values from a blockwise quantized tensor as + `scale * (input - zero_point)`. A missing `zero_points` operand means a + zero point of 0. + }]; + + let arguments = (ins + TensorOf<[TFL_I2, TFL_I4, I8, I16]>:$input, + TFL_TensorOf<[F8E8M0FNU, F16, BF16, F32]>:$scales, + TFL_TensorOfOrNone<[TFL_I2, TFL_I4, I8, I16, F32]>:$zero_points, + I64ArrayAttr:$block_shape, + BoolAttr:$symmetric + ); + + let results = (outs TFL_FpTensor:$output); + + let hasVerifier = 1; +} + def TFL_FakeQuantOp : TFL_Op<"fake_quant", [ Pure, QuantizableResult]> { diff --git a/tensorflow/compiler/mlir/lite/python/stablehlo_tfl_pipeline.cc b/tensorflow/compiler/mlir/lite/python/stablehlo_tfl_pipeline.cc index d91aed63f18f61..9518f575636555 100644 --- a/tensorflow/compiler/mlir/lite/python/stablehlo_tfl_pipeline.cc +++ b/tensorflow/compiler/mlir/lite/python/stablehlo_tfl_pipeline.cc @@ -150,6 +150,8 @@ void AddPipelinePasses(mlir::OpPassManager& pass_manager, // func.call boundaries that cause func.call result type mismatches when // quantized types are introduced. pass_manager.addPass(mlir::createInlinerPass()); + pass_manager.addNestedPass( + mlir::TFL::CreateFoldStablehloConstantTransformsPass()); pass_manager.addPass(mlir::TFL::CreateLowerQuantAnnotationsPass()); pass_manager.addPass(mlir::createSymbolDCEPass()); diff --git a/tensorflow/compiler/mlir/lite/schema/schema.fbs b/tensorflow/compiler/mlir/lite/schema/schema.fbs index b51dce02b63040..44af9948988d5b 100644 --- a/tensorflow/compiler/mlir/lite/schema/schema.fbs +++ b/tensorflow/compiler/mlir/lite/schema/schema.fbs @@ -1023,6 +1023,18 @@ table FullyConnectedOptions { // Parameters for FullyConnected version 11 or above. // When set, quantized_bias_type defines the dtype for both bias and accumulator. quantized_bias_type: TensorType; + + // Parameters for FullyConnected version 15 or above. + // + // An opaque, self-describing description of a non-standard quantization + // contract that the op's tensor-level quantization parameters cannot express + // on their own (for example a dynamically quantized, blockwise activation, or + // a weight grid that is centered between integer values). + // + // The payload is an encoded FlexBuffers map. A runtime that does not + // understand the contract named in the map must reject the op rather than + // fall back to standard quantization semantics. + quant_spec: [ubyte]; } table SoftmaxOptions { diff --git a/tensorflow/compiler/mlir/lite/schema/schema_generated.h b/tensorflow/compiler/mlir/lite/schema/schema_generated.h index cb9f9a479665f9..1bb5ef0f6a6fe6 100755 --- a/tensorflow/compiler/mlir/lite/schema/schema_generated.h +++ b/tensorflow/compiler/mlir/lite/schema/schema_generated.h @@ -9101,6 +9101,7 @@ struct FullyConnectedOptionsT : public ::flatbuffers::NativeTable { bool keep_num_dims = false; bool asymmetric_quantize_inputs = false; tflite::TensorType quantized_bias_type = tflite::TensorType_FLOAT32; + std::vector quant_spec{}; }; struct FullyConnectedOptions FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table { @@ -9111,7 +9112,8 @@ struct FullyConnectedOptions FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Ta VT_WEIGHTS_FORMAT = 6, VT_KEEP_NUM_DIMS = 8, VT_ASYMMETRIC_QUANTIZE_INPUTS = 10, - VT_QUANTIZED_BIAS_TYPE = 12 + VT_QUANTIZED_BIAS_TYPE = 12, + VT_QUANT_SPEC = 14 }; tflite::ActivationFunctionType fused_activation_function() const { return static_cast(GetField(VT_FUSED_ACTIVATION_FUNCTION, 0)); @@ -9128,6 +9130,9 @@ struct FullyConnectedOptions FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Ta tflite::TensorType quantized_bias_type() const { return static_cast(GetField(VT_QUANTIZED_BIAS_TYPE, 0)); } + const ::flatbuffers::Vector *quant_spec() const { + return GetPointer *>(VT_QUANT_SPEC); + } template bool Verify(::flatbuffers::VerifierTemplate &verifier) const { return VerifyTableStart(verifier) && @@ -9136,6 +9141,8 @@ struct FullyConnectedOptions FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Ta VerifyField(verifier, VT_KEEP_NUM_DIMS, 1) && VerifyField(verifier, VT_ASYMMETRIC_QUANTIZE_INPUTS, 1) && VerifyField(verifier, VT_QUANTIZED_BIAS_TYPE, 1) && + VerifyOffset(verifier, VT_QUANT_SPEC) && + verifier.VerifyVector(quant_spec()) && verifier.EndTable(); } FullyConnectedOptionsT *UnPack(const ::flatbuffers::resolver_function_t *_resolver = nullptr) const; @@ -9162,6 +9169,9 @@ struct FullyConnectedOptionsBuilder { void add_quantized_bias_type(tflite::TensorType quantized_bias_type) { fbb_.AddElement(FullyConnectedOptions::VT_QUANTIZED_BIAS_TYPE, static_cast(quantized_bias_type), 0); } + void add_quant_spec(::flatbuffers::Offset<::flatbuffers::Vector> quant_spec) { + fbb_.AddOffset(FullyConnectedOptions::VT_QUANT_SPEC, quant_spec); + } explicit FullyConnectedOptionsBuilder(::flatbuffers::FlatBufferBuilder &_fbb) : fbb_(_fbb) { start_ = fbb_.StartTable(); @@ -9179,8 +9189,10 @@ inline ::flatbuffers::Offset CreateFullyConnectedOptions( tflite::FullyConnectedOptionsWeightsFormat weights_format = tflite::FullyConnectedOptionsWeightsFormat_DEFAULT, bool keep_num_dims = false, bool asymmetric_quantize_inputs = false, - tflite::TensorType quantized_bias_type = tflite::TensorType_FLOAT32) { + tflite::TensorType quantized_bias_type = tflite::TensorType_FLOAT32, + ::flatbuffers::Offset<::flatbuffers::Vector> quant_spec = 0) { FullyConnectedOptionsBuilder builder_(_fbb); + builder_.add_quant_spec(quant_spec); builder_.add_quantized_bias_type(quantized_bias_type); builder_.add_asymmetric_quantize_inputs(asymmetric_quantize_inputs); builder_.add_keep_num_dims(keep_num_dims); @@ -9189,6 +9201,25 @@ inline ::flatbuffers::Offset CreateFullyConnectedOptions( return builder_.Finish(); } +inline ::flatbuffers::Offset CreateFullyConnectedOptionsDirect( + ::flatbuffers::FlatBufferBuilder &_fbb, + tflite::ActivationFunctionType fused_activation_function = tflite::ActivationFunctionType_NONE, + tflite::FullyConnectedOptionsWeightsFormat weights_format = tflite::FullyConnectedOptionsWeightsFormat_DEFAULT, + bool keep_num_dims = false, + bool asymmetric_quantize_inputs = false, + tflite::TensorType quantized_bias_type = tflite::TensorType_FLOAT32, + const std::vector *quant_spec = nullptr) { + auto quant_spec__ = quant_spec ? _fbb.CreateVector(*quant_spec) : 0; + return tflite::CreateFullyConnectedOptions( + _fbb, + fused_activation_function, + weights_format, + keep_num_dims, + asymmetric_quantize_inputs, + quantized_bias_type, + quant_spec__); +} + ::flatbuffers::Offset CreateFullyConnectedOptions(::flatbuffers::FlatBufferBuilder &_fbb, const FullyConnectedOptionsT *_o, const ::flatbuffers::rehasher_function_t *_rehasher = nullptr); struct SoftmaxOptionsT : public ::flatbuffers::NativeTable { @@ -18986,6 +19017,7 @@ inline void FullyConnectedOptions::UnPackTo(FullyConnectedOptionsT *_o, const :: { auto _e = keep_num_dims(); _o->keep_num_dims = _e; } { auto _e = asymmetric_quantize_inputs(); _o->asymmetric_quantize_inputs = _e; } { auto _e = quantized_bias_type(); _o->quantized_bias_type = _e; } + { auto _e = quant_spec(); if (_e) { _o->quant_spec.resize(_e->size()); std::copy(_e->begin(), _e->end(), _o->quant_spec.begin()); } } } inline ::flatbuffers::Offset CreateFullyConnectedOptions(::flatbuffers::FlatBufferBuilder &_fbb, const FullyConnectedOptionsT *_o, const ::flatbuffers::rehasher_function_t *_rehasher) { @@ -19001,13 +19033,15 @@ inline ::flatbuffers::Offset FullyConnectedOptions::Pack( auto _keep_num_dims = _o->keep_num_dims; auto _asymmetric_quantize_inputs = _o->asymmetric_quantize_inputs; auto _quantized_bias_type = _o->quantized_bias_type; + auto _quant_spec = _o->quant_spec.size() ? _fbb.CreateVector(_o->quant_spec) : 0; return tflite::CreateFullyConnectedOptions( _fbb, _fused_activation_function, _weights_format, _keep_num_dims, _asymmetric_quantize_inputs, - _quantized_bias_type); + _quantized_bias_type, + _quant_spec); } inline SoftmaxOptionsT *SoftmaxOptions::UnPack(const ::flatbuffers::resolver_function_t *_resolver) const { diff --git a/tensorflow/compiler/mlir/lite/tests/blockwise-ops.mlir b/tensorflow/compiler/mlir/lite/tests/blockwise-ops.mlir new file mode 100644 index 00000000000000..b4adc4d2a73926 --- /dev/null +++ b/tensorflow/compiler/mlir/lite/tests/blockwise-ops.mlir @@ -0,0 +1,126 @@ +// Copyright 2026 The TensorFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// ============================================================================== + +// Verifier tests for the blockwise quantize/dequantize ops. +// +// These two ops exist only between `LowerQuantAnnotations` and +// `FuseA4W2DynamicRangeFullyConnected`, so they are deliberately absent from +// the TFLite schema and are not run through `-tfl-runtime-verify`. + +// RUN: litert-opt -split-input-file -verify-diagnostics %s | FileCheck %s + +// CHECK-LABEL: testBlockwiseQuantize +func.func @testBlockwiseQuantize(%arg0: tensor<4x64xf32>) -> tensor<4x64xi4> { + // CHECK: "tfl.blockwise_quantize" + %0:3 = "tfl.blockwise_quantize"(%arg0) {block_shape = [1, 32], scale_type = f8E8M0FNU, symmetric = true, range_dilation = 1.500000e+00 : f32} : (tensor<4x64xf32>) -> (tensor<4x64xi4>, tensor<4x2xf8E8M0FNU>, none) + func.return %0#0 : tensor<4x64xi4> +} + +// ----- + +func.func @testBlockwiseQuantizeDynamicShape(%arg0: tensor) -> tensor { + // expected-error @+1 {{expects a statically shaped tensor}} + %0:3 = "tfl.blockwise_quantize"(%arg0) {block_shape = [1, 32], scale_type = f8E8M0FNU, symmetric = true} : (tensor) -> (tensor, tensor, none) + func.return %0#0 : tensor +} + +// ----- + +func.func @testBlockwiseQuantizeBlockShapeRank(%arg0: tensor<4x64xf32>) -> tensor<4x64xi4> { + // expected-error @+1 {{expects block_shape of rank 2}} + %0:3 = "tfl.blockwise_quantize"(%arg0) {block_shape = [32], scale_type = f8E8M0FNU, symmetric = true} : (tensor<4x64xf32>) -> (tensor<4x64xi4>, tensor<4x2xf8E8M0FNU>, none) + func.return %0#0 : tensor<4x64xi4> +} + +// ----- + +func.func @testBlockwiseQuantizeIndivisibleBlock(%arg0: tensor<4x64xf32>) -> tensor<4x64xi4> { + // expected-error @+1 {{expects dimension 1 (64) to be divisible by block_shape[1] (7)}} + %0:3 = "tfl.blockwise_quantize"(%arg0) {block_shape = [1, 7], scale_type = f8E8M0FNU, symmetric = true} : (tensor<4x64xf32>) -> (tensor<4x64xi4>, tensor<4x2xf8E8M0FNU>, none) + func.return %0#0 : tensor<4x64xi4> +} + +// ----- + +func.func @testBlockwiseQuantizeNonPositiveBlock(%arg0: tensor<4x64xf32>) -> tensor<4x64xi4> { + // expected-error @+1 {{expects a positive block_shape}} + %0:3 = "tfl.blockwise_quantize"(%arg0) {block_shape = [1, 0], scale_type = f8E8M0FNU, symmetric = true} : (tensor<4x64xf32>) -> (tensor<4x64xi4>, tensor<4x2xf8E8M0FNU>, none) + func.return %0#0 : tensor<4x64xi4> +} + +// ----- + +func.func @testBlockwiseQuantizeScaleGridMismatch(%arg0: tensor<4x64xf32>) -> tensor<4x64xi4> { + // expected-error @+1 {{expects scale dimension 1 to be 1 or 2}} + %0:3 = "tfl.blockwise_quantize"(%arg0) {block_shape = [1, 32], scale_type = f8E8M0FNU, symmetric = true} : (tensor<4x64xf32>) -> (tensor<4x64xi4>, tensor<4x8xf8E8M0FNU>, none) + func.return %0#0 : tensor<4x64xi4> +} + +// ----- + +func.func @testBlockwiseQuantizeScaleTypeMismatch(%arg0: tensor<4x64xf32>) -> tensor<4x64xi4> { + // expected-error @+1 {{expects the scale element type to match scale_type}} + %0:3 = "tfl.blockwise_quantize"(%arg0) {block_shape = [1, 32], scale_type = f8E8M0FNU, symmetric = true} : (tensor<4x64xf32>) -> (tensor<4x64xi4>, tensor<4x2xf32>, none) + func.return %0#0 : tensor<4x64xi4> +} + +// ----- + +// CHECK-LABEL: testBlockwiseDequantize +func.func @testBlockwiseDequantize(%arg0: tensor<32x64xi2>, %arg1: tensor<32x1xf32>, %arg2: tensor<32x1xf32>) -> tensor<32x64xf32> { + // CHECK: "tfl.blockwise_dequantize" + %0 = "tfl.blockwise_dequantize"(%arg0, %arg1, %arg2) {block_shape = [1, 64], symmetric = false} : (tensor<32x64xi2>, tensor<32x1xf32>, tensor<32x1xf32>) -> tensor<32x64xf32> + func.return %0 : tensor<32x64xf32> +} + +// ----- + +// A missing zero point is how a symmetric dequantize is spelled. +// CHECK-LABEL: testBlockwiseDequantizeNoZeroPoint +func.func @testBlockwiseDequantizeNoZeroPoint(%arg0: tensor<32x64xi2>, %arg1: tensor<32x1xf32>) -> tensor<32x64xf32> { + %none = "tfl.no_value"() {value} : () -> none + // CHECK: "tfl.blockwise_dequantize" + %0 = "tfl.blockwise_dequantize"(%arg0, %arg1, %none) {block_shape = [1, 64], symmetric = true} : (tensor<32x64xi2>, tensor<32x1xf32>, none) -> tensor<32x64xf32> + func.return %0 : tensor<32x64xf32> +} + +// ----- + +func.func @testBlockwiseDequantizeScalesRank(%arg0: tensor<32x64xi2>, %arg1: tensor<32xf32>) -> tensor<32x64xf32> { + %none = "tfl.no_value"() {value} : () -> none + // expected-error @+1 {{expects scales of rank 2}} + %0 = "tfl.blockwise_dequantize"(%arg0, %arg1, %none) {block_shape = [1, 64], symmetric = true} : (tensor<32x64xi2>, tensor<32xf32>, none) -> tensor<32x64xf32> + func.return %0 : tensor<32x64xf32> +} + +// ----- + +func.func @testBlockwiseDequantizeIndivisibleBlock(%arg0: tensor<32x64xi2>, %arg1: tensor<32x1xf32>) -> tensor<32x64xf32> { + %none = "tfl.no_value"() {value} : () -> none + // expected-error @+1 {{expects dimension 1 (64) to be divisible by block_shape[1] (7)}} + %0 = "tfl.blockwise_dequantize"(%arg0, %arg1, %none) {block_shape = [1, 7], symmetric = true} : (tensor<32x64xi2>, tensor<32x1xf32>, none) -> tensor<32x64xf32> + func.return %0 : tensor<32x64xf32> +} + +// ----- + +func.func @testBlockwiseDequantizeSymmetricWithZeroPoints(%arg0: tensor<32x64xi2>, %arg1: tensor<32x1xf32>) -> tensor<32x64xf32> { + %zp = "tfl.pseudo_const"() {value = dense<[[0.0], [-5.000000e-01]]> : tensor<2x1xf32>} : () -> tensor<2x1xf32> + %w = "tfl.pseudo_const"() {value = dense<1> : tensor<2x64xi2>} : () -> tensor<2x64xi2> + %s = "tfl.pseudo_const"() {value = dense<2.500000e-01> : tensor<2x1xf32>} : () -> tensor<2x1xf32> + // expected-error @+1 {{expects a per-tensor zero_point when symmetric is set}} + %0 = "tfl.blockwise_dequantize"(%w, %s, %zp) {block_shape = [1, 64], symmetric = true} : (tensor<2x64xi2>, tensor<2x1xf32>, tensor<2x1xf32>) -> tensor<2x64xf32> + func.return %arg0 : tensor<32x64xi2> +} diff --git a/tensorflow/compiler/mlir/lite/tools/versioning/op_version.cc b/tensorflow/compiler/mlir/lite/tools/versioning/op_version.cc index f4eedaa48f0b9a..d1bd505883c463 100644 --- a/tensorflow/compiler/mlir/lite/tools/versioning/op_version.cc +++ b/tensorflow/compiler/mlir/lite/tools/versioning/op_version.cc @@ -183,6 +183,15 @@ int GetBuiltinOperatorVersion(const OpSignature& op_sig) { reinterpret_cast(op_sig.builtin_data); TFLITE_DCHECK(fully_connected_params != nullptr); + // A non-empty `quant_spec` describes numerics that a runtime without + // explicit support for that spec cannot reproduce. It must be checked + // before the cases below, because an op carrying a `quant_spec` also + // matches them (an a4w2 op has int2 weights) and would otherwise be + // reported as a version that older runtimes accept. + if (fully_connected_params->quant_spec_size > 0) { + return 15; + } + if (op_sig.inputs.at(1).type == kTfLiteInt2) { return 14; } diff --git a/tensorflow/compiler/mlir/lite/tools/versioning/runtime_version.cc b/tensorflow/compiler/mlir/lite/tools/versioning/runtime_version.cc index 76074657c25fd4..9efeb19cfc2816 100644 --- a/tensorflow/compiler/mlir/lite/tools/versioning/runtime_version.cc +++ b/tensorflow/compiler/mlir/lite/tools/versioning/runtime_version.cc @@ -147,6 +147,7 @@ std::string FindMinimumRuntimeVersionForOp(tflite::BuiltinOperator op_code, {{BuiltinOperator_FULLY_CONNECTED, 12}, "2.17.0"}, {{BuiltinOperator_FULLY_CONNECTED, 13}, "2.18.0"}, {{BuiltinOperator_FULLY_CONNECTED, 14}, "2.21.0"}, + {{BuiltinOperator_FULLY_CONNECTED, 15}, "2.23.0"}, {{BuiltinOperator_GATHER, 1}, "1.6.0"}, {{BuiltinOperator_GATHER, 2}, "1.14.0"}, {{BuiltinOperator_GATHER, 3}, "1.15.0"}, diff --git a/tensorflow/compiler/mlir/lite/transforms/fold_stablehlo_constant_transforms_pass.cc b/tensorflow/compiler/mlir/lite/transforms/fold_stablehlo_constant_transforms_pass.cc new file mode 100644 index 00000000000000..7d0ce95d969d26 --- /dev/null +++ b/tensorflow/compiler/mlir/lite/transforms/fold_stablehlo_constant_transforms_pass.cc @@ -0,0 +1,277 @@ +/* Copyright 2026 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include +#include +#include +#include +#include + +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/FormatVariadic.h" +#include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project +#include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project +#include "mlir/IR/AsmState.h" // from @llvm-project +#include "mlir/IR/Attributes.h" // from @llvm-project +#include "mlir/IR/Builders.h" // from @llvm-project +#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project +#include "mlir/IR/BuiltinTypeInterfaces.h" // from @llvm-project +#include "mlir/IR/DialectResourceBlobManager.h" // from @llvm-project +#include "mlir/IR/MLIRContext.h" // from @llvm-project +#include "mlir/IR/Matchers.h" // from @llvm-project +#include "mlir/IR/PatternMatch.h" // from @llvm-project +#include "mlir/IR/Value.h" // from @llvm-project +#include "mlir/Pass/Pass.h" // from @llvm-project +#include "mlir/Support/LLVM.h" // from @llvm-project +#include "mlir/Support/LogicalResult.h" // from @llvm-project +#include "mlir/Support/TypeID.h" // from @llvm-project +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project +#include "stablehlo/dialect/StablehloOps.h" // from @stablehlo +#include "tensorflow/compiler/mlir/lite/transforms/passes.h" + +namespace mlir { +namespace TFL { +namespace { + +#define GEN_PASS_DEF_FOLDSTABLEHLOCONSTANTTRANSFORMSPASS +#include "tensorflow/compiler/mlir/lite/transforms/passes.h.inc" + +static void ComputePermutation( + llvm::ArrayRef perms, llvm::ArrayRef output_shape, + const char* raw_input, const int element_byte_size, + const int64_t current_axis, char*& output_element_addr, + llvm::MutableArrayRef current_input_index, + mlir::ShapedType input_shape_type) { + const int64_t input_axis = perms[current_axis]; + const bool is_last_axis = current_axis == output_shape.size() - 1; + + for (int i = 0; i < output_shape[current_axis]; ++i) { + current_input_index[input_axis] = i; + if (is_last_axis) { + int64_t input_flat_index = ElementsAttr::getFlattenedIndex( + input_shape_type, current_input_index); + const char* input_element_addr = + raw_input + input_flat_index * element_byte_size; + std::memcpy(output_element_addr, input_element_addr, element_byte_size); + output_element_addr += element_byte_size; + } else { + ComputePermutation(perms, output_shape, raw_input, element_byte_size, + current_axis + 1, output_element_addr, + current_input_index, input_shape_type); + } + } +} + +static ElementsAttr GetElementsAttrFromValue(mlir::Value val) { + if (!val) return nullptr; + ElementsAttr attr; + if (matchPattern(val, m_Constant(&attr))) { + return attr; + } + if (auto cst = val.getDefiningOp()) { + return mlir::dyn_cast_or_null(cst.getValue()); + } + if (auto cst = val.getDefiningOp()) { + return mlir::dyn_cast_or_null(cst.getValue()); + } + return nullptr; +} + +// Pattern to fold stablehlo.reshape on constant inputs. +struct FoldStablehloReshapeOpPattern + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(stablehlo::ReshapeOp op, + PatternRewriter& rewriter) const override { + mlir::Value input = op.getOperand(); + ElementsAttr input_attr = GetElementsAttrFromValue(input); + if (!input_attr) return failure(); + + auto result_type = mlir::cast(op.getType()); + if (!result_type.hasStaticShape()) return failure(); + + if (auto dense_elements = mlir::dyn_cast(input_attr)) { + rewriter.replaceOpWithNewOp( + op, dense_elements.reshape(result_type)); + return success(); + } else if (auto dense_resource = + mlir::dyn_cast(input_attr)) { + AsmResourceBlob* blob = dense_resource.getRawHandle().getBlob(); + if (!blob && dense_resource.getRawHandle().getResource()) { + blob = dense_resource.getRawHandle().getResource()->getBlob(); + } + if (!blob) return failure(); + + static std::atomic reshape_counter{0}; + auto key = dense_resource.getRawHandle().getKey(); + std::string new_key = + llvm::formatv("{0}_reshaped_{1}", key, reshape_counter++); + + auto new_blob = mlir::HeapAsmResourceBlob::allocate( + blob->getData().size(), /*align=*/64, true); + std::memcpy(const_cast(new_blob.getData().data()), + blob->getData().data(), blob->getData().size()); + auto new_attr = DenseResourceElementsAttr::get(result_type, new_key, + std::move(new_blob)); + rewriter.replaceOpWithNewOp(op, new_attr); + return success(); + } + return failure(); + } +}; + +// Pattern to fold stablehlo.transpose on constant inputs. +struct FoldStablehloTransposeOpPattern + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(stablehlo::TransposeOp op, + PatternRewriter& rewriter) const override { + mlir::Value input = op.getOperand(); + ElementsAttr input_attr = GetElementsAttrFromValue(input); + if (!input_attr) return failure(); + + auto input_type = mlir::cast(input_attr.getType()); + auto result_type = mlir::cast(op.getType()); + if (!result_type.hasStaticShape() || !input_type.hasStaticShape()) { + return failure(); + } + + if (!input_type.getElementType().isIntOrIndexOrFloat()) { + return failure(); + } + + const int num_dimensions = input_type.getRank(); + if (num_dimensions <= 1) { + rewriter.replaceOp(op, input); + return success(); + } + + llvm::ArrayRef perms = op.getPermutation(); + if (static_cast(perms.size()) != num_dimensions) { + return failure(); + } + + llvm::SmallVector output_shape; + for (int i = 0; i < num_dimensions; ++i) { + output_shape.push_back(input_type.getDimSize(perms[i])); + } + + if (auto dense_elements = mlir::dyn_cast(input_attr)) { + if (dense_elements.isSplat()) { + rewriter.replaceOpWithNewOp( + op, dense_elements.reshape(result_type)); + return success(); + } + + const int element_byte_size = + dense_elements.getElementType().getIntOrFloatBitWidth() / 8; + if (element_byte_size <= 0) return failure(); + + const char* raw_input = dense_elements.getRawData().data(); + llvm::SmallVector current_input_index(num_dimensions, 0); + llvm::SmallVector raw_output_arr( + dense_elements.getRawData().begin(), + dense_elements.getRawData().end()); + char* raw_output = raw_output_arr.data(); + + ComputePermutation(perms, output_shape, raw_input, element_byte_size, + /*current_axis=*/0, raw_output, current_input_index, + input_type); + + if (!DenseElementsAttr::isValidRawBuffer(result_type, raw_output_arr)) { + return failure(); + } + + auto new_attr = + DenseElementsAttr::getFromRawBuffer(result_type, raw_output_arr); + rewriter.replaceOpWithNewOp(op, new_attr); + return success(); + } else if (auto dense_resource = + mlir::dyn_cast(input_attr)) { + AsmResourceBlob* blob = dense_resource.getRawHandle().getBlob(); + if (!blob && dense_resource.getRawHandle().getResource()) { + blob = dense_resource.getRawHandle().getResource()->getBlob(); + } + if (!blob) return failure(); + + const int element_byte_size = + input_type.getElementType().getIntOrFloatBitWidth() / 8; + if (element_byte_size <= 0) return failure(); + + int64_t total_elements = result_type.getNumElements(); + size_t total_bytes = total_elements * element_byte_size; + + auto raw_output_blob = mlir::HeapAsmResourceBlob::allocate( + total_bytes, /*align=*/64, /*dataIsMutable=*/true); + llvm::ArrayRef data = raw_output_blob.getDataAs(); + llvm::MutableArrayRef raw_output_arr(const_cast(data.data()), + data.size()); + char* raw_output = raw_output_arr.data(); + const char* raw_input = blob->getData().data(); + if (!raw_input) return failure(); + + llvm::SmallVector current_input_index(num_dimensions, 0); + ComputePermutation(perms, output_shape, raw_input, element_byte_size, + /*current_axis=*/0, raw_output, current_input_index, + input_type); + + static std::atomic transpose_counter{0}; + std::string new_key = llvm::formatv("stablehlo_transpose_fold_result_{0}", + transpose_counter++); + auto new_attr = DenseResourceElementsAttr::get( + result_type, new_key, std::move(raw_output_blob)); + rewriter.replaceOpWithNewOp(op, new_attr); + return success(); + } + + return failure(); + } +}; + +struct FoldStablehloConstantTransformsPass + : public impl::FoldStablehloConstantTransformsPassBase< + FoldStablehloConstantTransformsPass> { + public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID( + FoldStablehloConstantTransformsPass) + + void runOnOperation() override { + func::FuncOp func = getOperation(); + MLIRContext* context = func.getContext(); + + RewritePatternSet patterns(context); + patterns + .add( + context); + + if (failed(applyPatternsGreedily(func, std::move(patterns)))) { + signalPassFailure(); + } + } +}; + +} // namespace + +std::unique_ptr> +CreateFoldStablehloConstantTransformsPass() { + return std::make_unique(); +} + +} // namespace TFL +} // namespace mlir diff --git a/tensorflow/compiler/mlir/lite/transforms/large_constant_fold_pass.cc b/tensorflow/compiler/mlir/lite/transforms/large_constant_fold_pass.cc index 9aae3a93606c55..b09ee0bb3006a4 100644 --- a/tensorflow/compiler/mlir/lite/transforms/large_constant_fold_pass.cc +++ b/tensorflow/compiler/mlir/lite/transforms/large_constant_fold_pass.cc @@ -26,6 +26,7 @@ limitations under the License. #include "Eigen/Core" // from @eigen_archive #include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallPtrSet.h" #include "llvm/ADT/SmallVector.h" #include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project @@ -34,12 +35,15 @@ limitations under the License. #include "mlir/IR/Builders.h" // from @llvm-project #include "mlir/IR/BuiltinAttributes.h" // from @llvm-project #include "mlir/IR/BuiltinOps.h" // from @llvm-project +#include "mlir/IR/MLIRContext.h" // from @llvm-project #include "mlir/IR/Matchers.h" // from @llvm-project +#include "mlir/IR/PatternMatch.h" // from @llvm-project #include "mlir/IR/Types.h" // from @llvm-project #include "mlir/IR/Value.h" // from @llvm-project #include "mlir/Pass/Pass.h" // from @llvm-project #include "mlir/Support/LLVM.h" // from @llvm-project #include "mlir/Support/LogicalResult.h" // from @llvm-project +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project #include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h" #include "tensorflow/compiler/mlir/lite/ir/tfl_ops_in_place_transpose.h" #include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.h" @@ -460,10 +464,167 @@ LogicalResult FoldResourceOpPattern(ModuleOp module) { return success(); } +struct PushTransposeThroughBlockwiseDequantize + : public OpRewritePattern { + explicit PushTransposeThroughBlockwiseDequantize(MLIRContext* ctx) + : OpRewritePattern(ctx) {} + + LogicalResult matchAndRewrite(TransposeOp op, + PatternRewriter& rewriter) const override { + auto deq = op.getOperand(0).getDefiningOp(); + if (!deq) return failure(); + + if (!deq.getOutput().hasOneUse()) return failure(); + + ElementsAttr input_attr; + if (!matchPattern(deq.getInput(), m_Constant(&input_attr))) { + if (auto const_op = deq.getInput().getDefiningOp()) { + input_attr = const_op.getValue(); + } else { + return failure(); + } + } + + ElementsAttr perm_attr; + if (!matchPattern(op.getOperand(1), m_Constant(&perm_attr))) { + if (auto const_op = op.getOperand(1).getDefiningOp()) { + perm_attr = const_op.getValue(); + } else { + return failure(); + } + } + auto int_perm_attr = mlir::dyn_cast(perm_attr); + if (!int_perm_attr) return failure(); + + auto input_type = + mlir::dyn_cast(deq.getInput().getType()); + auto output_type = mlir::dyn_cast(op.getType()); + if (!input_type || !output_type) return failure(); + + int rank = input_type.getRank(); + auto perms = llvm::to_vector<4>( + llvm::map_range(int_perm_attr.getValues(), + [](const APInt& val) { return val.getSExtValue(); })); + if (perms.size() != rank) return failure(); + + llvm::SmallDenseSet seen_dims; + for (int64_t p : perms) { + if (p < 0 || p >= rank || !seen_dims.insert(p).second) return failure(); + } + + ArrayAttr new_block_shape_attr = nullptr; + if (auto block_shape_attr = deq.getBlockShapeAttr()) { + if (block_shape_attr.size() != rank) return failure(); + SmallVector new_block_shape(rank); + for (int i = 0; i < rank; ++i) { + new_block_shape[i] = + mlir::cast(block_shape_attr[perms[i]]).getInt(); + } + new_block_shape_attr = rewriter.getI64ArrayAttr(new_block_shape); + } else { + return failure(); + } + + SmallVector transposed_input_shape; + for (int64_t p : perms) { + transposed_input_shape.push_back(input_type.getDimSize(p)); + } + auto transposed_input_type = RankedTensorType::get( + transposed_input_shape, input_type.getElementType()); + Value new_input = rewriter.create( + op.getLoc(), transposed_input_type, deq.getInput(), op.getOperand(1)); + + Value new_scales = deq.getScales(); + if (new_scales && + !(new_scales.getDefiningOp() && + mlir::isa(new_scales.getDefiningOp()))) { + auto scales_type = mlir::dyn_cast(new_scales.getType()); + if (!scales_type || scales_type.getRank() != rank) { + return failure(); + } + SmallVector transposed_scales_shape; + for (int64_t p : perms) { + transposed_scales_shape.push_back(scales_type.getDimSize(p)); + } + auto transposed_scales_type = RankedTensorType::get( + transposed_scales_shape, scales_type.getElementType()); + ElementsAttr scales_attr; + if (matchPattern(new_scales, m_Constant(&scales_attr))) { + if (auto dense_attr = mlir::dyn_cast(scales_attr)) { + if (dense_attr.isSplat()) { + new_scales = rewriter.create( + op.getLoc(), dense_attr.reshape(transposed_scales_type)); + } else { + new_scales = rewriter.create( + op.getLoc(), transposed_scales_type, new_scales, + op.getOperand(1)); + } + } else { + new_scales = + rewriter.create(op.getLoc(), transposed_scales_type, + new_scales, op.getOperand(1)); + } + } else { + new_scales = rewriter.create( + op.getLoc(), transposed_scales_type, new_scales, op.getOperand(1)); + } + } + + Value new_zp = deq.getZeroPoints(); + if (new_zp && !(new_zp.getDefiningOp() && + mlir::isa(new_zp.getDefiningOp()))) { + auto zp_type = mlir::dyn_cast(new_zp.getType()); + if (!zp_type || zp_type.getRank() != rank) { + return failure(); + } + SmallVector transposed_zp_shape; + for (int64_t p : perms) { + transposed_zp_shape.push_back(zp_type.getDimSize(p)); + } + auto transposed_zp_type = + RankedTensorType::get(transposed_zp_shape, zp_type.getElementType()); + ElementsAttr zp_attr; + if (matchPattern(new_zp, m_Constant(&zp_attr))) { + if (auto dense_attr = mlir::dyn_cast(zp_attr)) { + if (dense_attr.isSplat()) { + new_zp = rewriter.create( + op.getLoc(), dense_attr.reshape(transposed_zp_type)); + } else { + new_zp = rewriter.create( + op.getLoc(), transposed_zp_type, new_zp, op.getOperand(1)); + } + } else { + new_zp = rewriter.create(op.getLoc(), transposed_zp_type, + new_zp, op.getOperand(1)); + } + } else { + new_zp = rewriter.create(op.getLoc(), transposed_zp_type, + new_zp, op.getOperand(1)); + } + } + + auto new_deq = rewriter.create( + op.getLoc(), output_type, new_input, new_scales, new_zp, + new_block_shape_attr, deq.getSymmetricAttr()); + rewriter.replaceOp(op, new_deq.getOutput()); + return success(); + } +}; + } // namespace void LargeConstantFoldPass::runOnOperation() { ModuleOp module = getOperation(); + MLIRContext* ctx = &getContext(); + + RewritePatternSet patterns(ctx); + patterns.add(ctx); + GreedyRewriteConfig config; + config.enableFolding(false); + if (failed(applyPatternsGreedily(module, std::move(patterns), config))) { + signalPassFailure(); + return; + } if (failed(FoldResourceOpPattern(module))) { signalPassFailure(); } diff --git a/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.h b/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.h index a7ebd380742df7..a396accce5c515 100644 --- a/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.h +++ b/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.h @@ -18,9 +18,12 @@ limitations under the License. #include +#include "llvm/ADT/StringRef.h" #include "llvm/Support/Casting.h" +#include "mlir/IR/Attributes.h" // from @llvm-project #include "mlir/IR/Builders.h" // from @llvm-project #include "mlir/IR/BuiltinAttributes.h" // from @llvm-project +#include "mlir/IR/BuiltinTypes.h" // from @llvm-project #include "mlir/IR/Location.h" // from @llvm-project #include "mlir/IR/MLIRContext.h" // from @llvm-project #include "mlir/IR/Types.h" // from @llvm-project @@ -30,106 +33,208 @@ limitations under the License. namespace mlir::TFL { -template -LogicalResult FillCompositeParams(CompositeOpType op, - SmallVector& scales, - SmallVector& zero_points, - int& num_bits, bool& is_signed, - bool& is_narrow_range) { - auto scale_attr = llvm::dyn_cast_or_null( - op.getCompositeAttributes().get("scale")); - if (scale_attr == nullptr) { - return failure(); - } - for (auto float_attr : scale_attr.template getValues()) { - scales.push_back(float_attr.getValue().convertToDouble()); - } +// Quantization annotations reach this pass as either a `stablehlo.composite` or +// a `stablehlo.custom_call`. The two carry the same information but spell the +// accessors differently, so these overloads let the templates below handle both +// without duplicating the extraction logic. +inline Attribute GetOpAttribute(stablehlo::CompositeOp op, StringRef name) { + return op.getCompositeAttributes().get(name); +} - auto zero_point_attr = llvm::dyn_cast_or_null( - op.getCompositeAttributes().get("zero_point")); - if (zero_point_attr == nullptr) { - for (int i = 0; i < scales.size(); ++i) { - zero_points.push_back(0); - } - } else if (zero_point_attr.isSplat()) { - for (int i = 0; i < scales.size(); ++i) { - zero_points.push_back( - zero_point_attr.template getSplatValue().getInt()); - } - } else { - for (IntegerAttr zp : zero_point_attr.template getValues()) { - zero_points.push_back(zp.getInt()); - } - } +inline Attribute GetOpAttribute(stablehlo::CustomCallOp op, StringRef name) { + return op->getAttr(name); +} - auto dtype_attr = llvm::dyn_cast_or_null( - op.getCompositeAttributes().get("dtype")); - if (dtype_attr == nullptr) { - return failure(); +inline StringRef GetQuantAnnotationName(stablehlo::CompositeOp op) { + return op.getName(); +} + +inline StringRef GetQuantAnnotationName(stablehlo::CustomCallOp op) { + return op.getCallTargetName(); +} + +// The integer storage of a quantized tensor, as named by the annotation's +// `dtype` attribute. +struct QuantStorage { + Type type; + int num_bits = 0; + bool is_signed = false; + + // Inclusive bounds of the representable range. `narrow_range` drops the most + // negative value so that the range is symmetric around zero. + // + // This is the single definition of the storage range used by the whole + // lowering: the constant folder that quantizes weights at compile time and + // the runtime kernels must agree on it exactly, or the model's numerics + // change depending on which one ran. + int64_t Min(bool narrow_range) const { + if (!is_signed) return 0; + const int64_t min = -(int64_t{1} << (num_bits - 1)); + return narrow_range ? min + 1 : min; + } + int64_t Max(bool narrow_range) const { + if (!is_signed) return (int64_t{1} << num_bits) - 1; + return (int64_t{1} << (num_bits - 1)) - 1; } - auto dtype = dtype_attr.getValue(); +}; +// Parses the annotation's `dtype` string. Fails on an unrecognized name rather +// than defaulting, because every caller uses the result to decide how to round +// and clamp real values. +inline LogicalResult ParseQuantStorage(StringRef dtype, MLIRContext* ctx, + QuantStorage& storage) { if (dtype == "i2") { - num_bits = 2; - is_signed = true; + storage = {IntegerType::get(ctx, 2), 2, /*is_signed=*/true}; } else if (dtype == "i4") { - num_bits = 4; - is_signed = true; + storage = {IntegerType::get(ctx, 4), 4, /*is_signed=*/true}; } else if (dtype == "ui4") { - num_bits = 4; - is_signed = false; + storage = {IntegerType::get(ctx, 4), 4, /*is_signed=*/false}; } else if (dtype == "i8") { - num_bits = 8; - is_signed = true; + storage = {IntegerType::get(ctx, 8), 8, /*is_signed=*/true}; + } else if (dtype == "ui8") { + storage = {IntegerType::get(ctx, 8), 8, /*is_signed=*/false}; } else if (dtype == "i16") { - num_bits = 16; - is_signed = true; + storage = {IntegerType::get(ctx, 16), 16, /*is_signed=*/true}; } else { return failure(); } - auto narrow_range_attr = llvm::dyn_cast_or_null( - op.getCompositeAttributes().get("narrow_range")); - if (narrow_range_attr == nullptr) { + return success(); +} + +// Everything a `quant.*` annotation says about a blockwise quantized tensor. +// +// All quantization parameters are annotation attributes rather than operands: +// the annotation op always has exactly one operand, the value being quantized. +struct BlockwiseAnnotation { + SmallVector block_shape; + QuantStorage storage; + bool narrow_range = false; + bool symmetric = false; + // Element type of the scale. `f8E8M0FNU` constrains the scale to a power of + // two; otherwise `f32`. + Type scale_type; + // Widens the quantized range used to derive a dynamic scale, so that + // outliers clip rather than stretching the whole block. 0 when absent. + double range_dilation = 0.0; + // Null for dynamic range quantization, where the scale is derived from the + // tensor at runtime. + ElementsAttr scale; + // Null when the quantization is symmetric. + ElementsAttr zero_point; + + bool IsDynamic() const { return scale == nullptr; } +}; + +// Returns true if `op` carries a `block_shape` attribute, i.e. it describes +// blockwise rather than per-tensor or per-axis quantization. +template +bool IsBlockwiseAnnotation(OpType op) { + return llvm::isa_and_nonnull( + GetOpAttribute(op, "block_shape")); +} + +// Extracts the blockwise quantization parameters of `op`. Only call this when +// `IsBlockwiseAnnotation(op)` holds. +// +// Fails on any attribute that is present but unrecognized. Falling back to a +// default would produce a model that runs but computes something other than +// what the annotation asked for. +template +LogicalResult ParseBlockwiseAnnotation(OpType op, BlockwiseAnnotation& out) { + MLIRContext* ctx = op->getContext(); + + auto block_shape_attr = llvm::dyn_cast_or_null( + GetOpAttribute(op, "block_shape")); + if (block_shape_attr == nullptr) return failure(); + for (const IntegerAttr block : + block_shape_attr.template getValues()) { + out.block_shape.push_back(block.getInt()); + } + + auto dtype_attr = + llvm::dyn_cast_or_null(GetOpAttribute(op, "dtype")); + if (dtype_attr == nullptr) return failure(); + if (failed(ParseQuantStorage(dtype_attr.getValue(), ctx, out.storage))) { return failure(); } - is_narrow_range = narrow_range_attr.getValue(); - return success(); -} + if (auto narrow_range = llvm::dyn_cast_or_null( + GetOpAttribute(op, "narrow_range"))) { + out.narrow_range = narrow_range.getValue(); + } + if (auto symmetric = + llvm::dyn_cast_or_null(GetOpAttribute(op, "symmetric"))) { + out.symmetric = symmetric.getValue(); + } -template -bool IsDrqFakeQuant(CompositeOpType op) { - if (op.getName() != "quant.fake_quant") { - return false; + out.scale_type = Float32Type::get(ctx); + if (Attribute scale_dtype = GetOpAttribute(op, "act_scale_dtype")) { + auto scale_dtype_str = llvm::dyn_cast(scale_dtype); + if (scale_dtype_str == nullptr) return failure(); + if (scale_dtype_str.getValue() == "e8m0") { + out.scale_type = Float8E8M0FNUType::get(ctx); + } else if (scale_dtype_str.getValue() != "fp32" && + scale_dtype_str.getValue() != "float32") { + return failure(); + } } - SmallVector scales; - SmallVector zero_points; - int num_bits; - bool is_signed; - bool is_narrow_range; - if (failed(FillCompositeParams(op, scales, zero_points, num_bits, is_signed, - is_narrow_range))) { - return false; + + if (Attribute dilation = GetOpAttribute(op, "range_dilation")) { + auto dilation_fp = llvm::dyn_cast(dilation); + if (dilation_fp == nullptr) return failure(); + out.range_dilation = dilation_fp.getValueAsDouble(); } - return scales.empty() && zero_points.empty(); + + // An empty `scale` means dynamic range quantization; see `IsDrqFakeQuant`. + auto scale = + llvm::dyn_cast_or_null(GetOpAttribute(op, "scale")); + if (scale != nullptr && !scale.empty()) out.scale = scale; + + auto zero_point = + llvm::dyn_cast_or_null(GetOpAttribute(op, "zero_point")); + if (zero_point != nullptr && !zero_point.empty()) out.zero_point = zero_point; + + return success(); +} + +// Returns true if `op` carries the static `scale` attribute that describes its +// quantization parameters, i.e. the parameters are baked into the annotation. +// +// An annotation without one describes dynamic range quantization (DRQ), where +// the parameters are derived from the tensor at runtime. This is the exact +// condition under which `FillCompositeParams` below fails, so the two must stay +// in sync. +template +bool HasStaticQuantParams(OpType op) { + auto scale_attr = + llvm::dyn_cast_or_null(GetOpAttribute(op, "scale")); + return scale_attr != nullptr && !scale_attr.empty(); } -template <> -inline LogicalResult FillCompositeParams( - stablehlo::CustomCallOp op, SmallVector& scales, - SmallVector& zero_points, int& num_bits, bool& is_signed, - bool& is_narrow_range) { +// Extracts the quantization parameters from a per-tensor or per-axis `quant.*` +// annotation into `scales`, `zero_points`, `num_bits`, `is_signed` and +// `is_narrow_range`. +// +// Fails if any required attribute is missing or unrecognized. In particular it +// fails on DRQ annotations, which carry no static `scale`; callers that need to +// handle those must check `IsDrqFakeQuant`/`HasStaticQuantParams` first. On +// success `scales` is guaranteed non-empty and `zero_points` has the same size. +template +LogicalResult FillCompositeParams(OpType op, SmallVector& scales, + SmallVector& zero_points, + int& num_bits, bool& is_signed, + bool& is_narrow_range) { auto scale_attr = - llvm::dyn_cast_or_null(op->getAttr("scale")); - if (scale_attr == nullptr) { + llvm::dyn_cast_or_null(GetOpAttribute(op, "scale")); + if (scale_attr == nullptr || scale_attr.empty()) { return failure(); } - for (auto float_attr : scale_attr.getValues()) { + for (auto float_attr : scale_attr.template getValues()) { scales.push_back(float_attr.getValue().convertToDouble()); } - auto zero_point_attr = - llvm::dyn_cast_or_null(op->getAttr("zero_point")); + auto zero_point_attr = llvm::dyn_cast_or_null( + GetOpAttribute(op, "zero_point")); if (zero_point_attr == nullptr) { for (int i = 0; i < scales.size(); ++i) { zero_points.push_back(0); @@ -137,40 +242,29 @@ inline LogicalResult FillCompositeParams( } else if (zero_point_attr.isSplat()) { for (int i = 0; i < scales.size(); ++i) { zero_points.push_back( - zero_point_attr.getSplatValue().getInt()); + zero_point_attr.template getSplatValue().getInt()); } } else { - for (IntegerAttr zp : zero_point_attr.getValues()) { + for (IntegerAttr zp : zero_point_attr.template getValues()) { zero_points.push_back(zp.getInt()); } } - auto dtype_attr = llvm::dyn_cast_or_null(op->getAttr("dtype")); + auto dtype_attr = + llvm::dyn_cast_or_null(GetOpAttribute(op, "dtype")); if (dtype_attr == nullptr) { return failure(); } - auto dtype = dtype_attr.getValue(); - - if (dtype == "i2") { - num_bits = 2; - is_signed = true; - } else if (dtype == "i4") { - num_bits = 4; - is_signed = true; - } else if (dtype == "ui4") { - num_bits = 4; - is_signed = false; - } else if (dtype == "i8") { - num_bits = 8; - is_signed = true; - } else if (dtype == "i16") { - num_bits = 16; - is_signed = true; - } else { + QuantStorage storage; + if (failed(ParseQuantStorage(dtype_attr.getValue(), op->getContext(), + storage))) { return failure(); } + num_bits = storage.num_bits; + is_signed = storage.is_signed; + auto narrow_range_attr = - llvm::dyn_cast_or_null(op->getAttr("narrow_range")); + llvm::dyn_cast_or_null(GetOpAttribute(op, "narrow_range")); if (narrow_range_attr == nullptr) { return failure(); } @@ -179,22 +273,12 @@ inline LogicalResult FillCompositeParams( return success(); } -template <> -inline bool IsDrqFakeQuant( - stablehlo::CustomCallOp op) { - if (op.getCallTargetName() != "quant.fake_quant") { - return false; - } - SmallVector scales; - SmallVector zero_points; - int num_bits; - bool is_signed; - bool is_narrow_range; - if (failed(FillCompositeParams(op, scales, zero_points, num_bits, is_signed, - is_narrow_range))) { +template +bool IsDrqFakeQuant(OpType op) { + if (GetQuantAnnotationName(op) != "quant.fake_quant") { return false; } - return scales.empty() && zero_points.empty(); + return !HasStaticQuantParams(op); } LogicalResult GetStorageParams(unsigned num_bits, bool narrow_range, diff --git a/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_pass.cc b/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_pass.cc index 9ecf7dca971e65..9eb737d5f71d70 100644 --- a/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_pass.cc +++ b/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_pass.cc @@ -15,10 +15,14 @@ limitations under the License. // This transformation pass applies quantization on TFLite dialect. +#include +#include +#include #include #include #include +#include "llvm/ADT/STLExtras.h" #include "llvm/Support/Casting.h" #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project #include "mlir/Dialect/Quant/IR/Quant.h" // from @llvm-project // IWYU pragma: keep @@ -28,6 +32,7 @@ limitations under the License. #include "mlir/IR/BuiltinTypes.h" // from @llvm-project #include "mlir/IR/Diagnostics.h" // from @llvm-project #include "mlir/IR/DialectResourceBlobManager.h" // from @llvm-project // IWYU pragma: keep +#include "mlir/IR/Operation.h" // from @llvm-project #include "mlir/IR/PatternMatch.h" // from @llvm-project #include "mlir/IR/Value.h" // from @llvm-project #include "mlir/Support/LLVM.h" // from @llvm-project @@ -49,6 +54,284 @@ namespace { #define GEN_PASS_DEF_LOWERQUANTANNOTATIONSPASS #include "tensorflow/compiler/mlir/lite/transforms/passes.h.inc" +//===----------------------------------------------------------------------===// +// Blockwise annotation lowering +//===----------------------------------------------------------------------===// + +// Returns the blob backing `attr`, or null if it has none. +AsmResourceBlob* GetResourceBlob(DenseResourceElementsAttr attr) { + if (AsmResourceBlob* blob = attr.getRawHandle().getBlob()) return blob; + if (auto* resource = attr.getRawHandle().getResource()) { + return resource->getBlob(); + } + return nullptr; +} + +// Appends the elements of a resource blob to `out`, reinterpreting the raw +// bytes according to `element_type`. +// +// A resource attribute has no per-element accessor, so the element type has to +// be dispatched on explicitly. Reading an unexpected type as `float` would +// produce plausible-looking garbage, so anything unhandled fails. +template +void AppendBlobAs(AsmResourceBlob* blob, SmallVectorImpl& out) { + for (T value : blob->getDataAs()) + out.push_back(static_cast(value)); +} + +LogicalResult ReadNumericElements(ElementsAttr attr, + SmallVectorImpl& out) { + if (auto dense_fp = mlir::dyn_cast(attr)) { + for (const APFloat& value : dense_fp.getValues()) { + out.push_back(value.convertToDouble()); + } + return success(); + } + if (auto dense_int = mlir::dyn_cast(attr)) { + for (const APInt& value : dense_int.getValues()) { + out.push_back(static_cast(value.getSExtValue())); + } + return success(); + } + + auto dense_resource = mlir::dyn_cast(attr); + if (dense_resource == nullptr) return failure(); + AsmResourceBlob* blob = GetResourceBlob(dense_resource); + if (blob == nullptr) return failure(); + + const Type element_type = dense_resource.getType().getElementType(); + if (element_type.isF32()) { + AppendBlobAs(blob, out); + } else if (element_type.isF64()) { + AppendBlobAs(blob, out); + } else if (element_type.isInteger(8)) { + AppendBlobAs(blob, out); + } else if (element_type.isInteger(16)) { + AppendBlobAs(blob, out); + } else if (element_type.isInteger(32)) { + AppendBlobAs(blob, out); + } else if (element_type.isInteger(64)) { + AppendBlobAs(blob, out); + } else { + return failure(); + } + return success(); +} + +ElementsAttr GetElementsAttrFromValue(Value value) { + if (!value) return nullptr; + mlir::Operation* defining_op = value.getDefiningOp(); + if (defining_op == nullptr) return nullptr; + if (auto reshape = mlir::dyn_cast(defining_op)) { + return GetElementsAttrFromValue(reshape.getOperand()); + } + if (auto reshape = mlir::dyn_cast(defining_op)) { + return GetElementsAttrFromValue(reshape.getInput()); + } + return defining_op->getAttrOfType("value"); +} + +// Maps each element of a tensor of shape `shape` to its block in a grid tiled +// by `block_shape`, and from there to an index into a parameter tensor of shape +// `param_shape`. A parameter dimension of size 1 broadcasts across that axis. +class BlockIndexer { + public: + BlockIndexer(ArrayRef shape, ArrayRef block_shape, + ArrayRef param_shape) + : shape_(shape), block_shape_(block_shape), param_shape_(param_shape) {} + + int64_t ParamIndexOf(int64_t flat_index) const { + int64_t param_index = 0; + int64_t remainder = flat_index; + // Recover the coordinates from the innermost dimension outwards, then fold + // them into the parameter index from the outermost inwards. + SmallVector block_coords(shape_.size()); + for (int dim = shape_.size() - 1; dim >= 0; --dim) { + block_coords[dim] = (remainder % shape_[dim]) / block_shape_[dim]; + remainder /= shape_[dim]; + } + for (size_t dim = 0; dim < shape_.size(); ++dim) { + const int64_t param_dim = param_shape_[dim]; + param_index = + param_index * param_dim + (param_dim == 1 ? 0 : block_coords[dim]); + } + return param_index; + } + + private: + ArrayRef shape_; + ArrayRef block_shape_; + ArrayRef param_shape_; +}; + +// Quantizes a constant at compile time using the blockwise parameters of +// `annotation`, mirroring `tfl.blockwise_dequantize`: +// +// q = clamp(nearbyint(real / scale + zero_point), q_min, q_max) +// +// `nearbyint` rounds halves to even, matching `jnp.round` on the JAX side. The +// clamp bounds come from `QuantStorage`, which the runtime kernels share. +FailureOr QuantizeBlockwise(ElementsAttr real_values, + const BlockwiseAnnotation& annotation, + Location loc) { + auto real_type = mlir::dyn_cast(real_values.getType()); + if (real_type == nullptr || !real_type.hasStaticShape()) { + return emitError(loc, "blockwise weight must have a static shape"); + } + const ArrayRef shape = real_type.getShape(); + if (annotation.block_shape.size() != shape.size()) { + return emitError(loc, "block_shape rank does not match the weight rank"); + } + for (auto [dim, dim_size] : llvm::enumerate(shape)) { + const int64_t block_size = annotation.block_shape[dim]; + if (block_size <= 0 || dim_size % block_size != 0) { + return emitError(loc, "weight dimension ") + << dim << " (" << dim_size << ") is not a multiple of block_shape[" + << dim << "] (" << block_size << ")"; + } + } + + SmallVector values; + SmallVector scales; + if (failed(ReadNumericElements(real_values, values)) || + failed(ReadNumericElements(annotation.scale, scales))) { + return emitError(loc, "unsupported element type in a blockwise weight"); + } + + SmallVector zero_points; + if (annotation.zero_point != nullptr && + failed(ReadNumericElements(annotation.zero_point, zero_points))) { + return emitError(loc, "unsupported element type in a blockwise zero point"); + } + + const BlockIndexer scale_indexer(shape, annotation.block_shape, + annotation.scale.getShapedType().getShape()); + const BlockIndexer zero_point_indexer( + shape, annotation.block_shape, + zero_points.empty() ? ArrayRef() + : annotation.zero_point.getShapedType().getShape()); + + const double q_min = annotation.storage.Min(annotation.narrow_range); + const double q_max = annotation.storage.Max(annotation.narrow_range); + + SmallVector quantized; + quantized.reserve(values.size()); + for (auto [index, value] : llvm::enumerate(values)) { + const int64_t scale_index = scale_indexer.ParamIndexOf(index); + if (scale_index >= static_cast(scales.size())) { + return emitError(loc, "blockwise scale is smaller than its block grid"); + } + const double scale = scales[scale_index]; + if (scale == 0.0) { + return emitError(loc, "blockwise scale contains a zero"); + } + + double zero_point = 0.0; + if (!zero_points.empty()) { + const int64_t index_of = zero_point_indexer.ParamIndexOf(index); + if (index_of >= static_cast(zero_points.size())) { + return emitError(loc, + "blockwise zero point is smaller than its block grid"); + } + zero_point = zero_points[index_of]; + } + + const double q = + std::clamp(std::nearbyint(value / scale + zero_point), q_min, q_max); + quantized.emplace_back(annotation.storage.num_bits, + static_cast(static_cast(q)), + /*isSigned=*/annotation.storage.is_signed); + } + + return cast(DenseElementsAttr::get( + RankedTensorType::get(shape, annotation.storage.type), quantized)); +} + +// Lowers a blockwise `quant.fake_quant` annotation. +// +// Weights carry a static scale and are quantized here, at compile time, into a +// constant feeding a `tfl.blockwise_dequantize`. Activations carry no scale: +// they get a `tfl.blockwise_quantize` that derives one at runtime, feeding the +// matching dequantize. +// +// Both shapes are collapsed back into a single op later, by +// `FuseA4W2DynamicRangeFullyConnected`. +template +LogicalResult LowerBlockwiseFakeQuant(OpType op, PatternRewriter& rewriter) { + BlockwiseAnnotation annotation; + if (failed(ParseBlockwiseAnnotation(op, annotation))) { + return op.emitError( + "blockwise quant.fake_quant has missing or unsupported attributes"); + } + + // The input is the last operand so that dynamically shaped models, which + // prepend shape operands, keep working. + Value input = op.getOperand(op.getNumOperands() - 1); + auto input_type = cast(input.getType()); + auto output_type = cast(op.getType(0)); + const ArrayAttr block_shape = + rewriter.getI64ArrayAttr(annotation.block_shape); + + Value quantized; + Value scale; + Value zero_point; + + if (annotation.IsDynamic()) { + SmallVector grid_shape; + grid_shape.reserve(input_type.getRank()); + for (auto [dim, dim_size] : llvm::enumerate(input_type.getShape())) { + grid_shape.push_back(dim_size / annotation.block_shape[dim]); + } + // A symmetric quantizer has no zero point to report; the op returns none + // rather than a grid of zeros so that the consumer cannot mistake it for + // a meaningful value. + const SmallVector zero_point_shape(input_type.getRank(), 1); + + auto blockwise_quantize = rewriter.create( + op.getLoc(), + /*output=*/ + RankedTensorType::get(input_type.getShape(), annotation.storage.type), + /*scale=*/RankedTensorType::get(grid_shape, annotation.scale_type), + /*zero_point=*/ + annotation.symmetric ? Type(rewriter.getNoneType()) + : Type(RankedTensorType::get( + zero_point_shape, annotation.storage.type)), + input, block_shape, TypeAttr::get(annotation.scale_type), + rewriter.getBoolAttr(annotation.symmetric), + rewriter.getF32FloatAttr(annotation.range_dilation)); + + quantized = blockwise_quantize.getOutput(); + scale = blockwise_quantize.getScale(); + zero_point = blockwise_quantize.getZeroPoint(); + } else { + ElementsAttr real_values = GetElementsAttrFromValue(input); + if (real_values == nullptr) { + return op.emitError( + "blockwise quant.fake_quant has a static scale but a non-constant " + "input"); + } + FailureOr quantized_values = + QuantizeBlockwise(real_values, annotation, op.getLoc()); + if (failed(quantized_values)) return failure(); + + quantized = rewriter.create(op.getLoc(), *quantized_values); + scale = rewriter.create(op.getLoc(), annotation.scale); + if (annotation.zero_point == nullptr) { + zero_point = rewriter.create( + op.getLoc(), rewriter.getNoneType(), rewriter.getUnitAttr()); + } else { + zero_point = + rewriter.create(op.getLoc(), annotation.zero_point); + } + } + + auto blockwise_dequantize = rewriter.create( + op.getLoc(), output_type, quantized, scale, zero_point, block_shape, + rewriter.getBoolAttr(annotation.symmetric)); + rewriter.replaceOp(op, blockwise_dequantize.getOutput()); + return success(); +} + /** * Replaces a quant.quantize composite op with a TFLite quantize op that outputs * a quantized tensor based on the composite op's attributes. @@ -298,7 +581,13 @@ class RewriteFakeQuantCompositeOp public: LogicalResult matchAndRewrite(stablehlo::CompositeOp op, PatternRewriter& rewriter) const final { - if (op.getName() != "quant.fake_quant" || IsDrqFakeQuant(op)) { + if (op.getName() != "quant.fake_quant") { + return failure(); + } + if (IsBlockwiseAnnotation(op)) { + return LowerBlockwiseFakeQuant(op, rewriter); + } + if (IsDrqFakeQuant(op)) { return failure(); } @@ -645,7 +934,13 @@ void LowerQuantAnnotationsPass::runOnOperation() { public: LogicalResult matchAndRewrite(stablehlo::CustomCallOp op, PatternRewriter& rewriter) const final { - if (op.getCallTargetName() != "quant.fake_quant" || IsDrqFakeQuant(op)) { + if (op.getCallTargetName() != "quant.fake_quant") { + return failure(); + } + if (IsBlockwiseAnnotation(op)) { + return LowerBlockwiseFakeQuant(op, rewriter); + } + if (IsDrqFakeQuant(op)) { return failure(); } diff --git a/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc b/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc index f90ba3d574d314..c8e4ce5c6077d9 100644 --- a/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc +++ b/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc @@ -53,7 +53,7 @@ bool IsFromFoldableChain(mlir::Value value) { while (defining_op) { // `QConstOp` carries its value in an attribute but is neither ConstantLike // nor foldable, so it has to be named rather than left to `m_Constant`. - if (mlir::isa(defining_op) || + if (mlir::isa(defining_op) || defining_op->hasTrait() || matchPattern(defining_op, mlir::m_Constant())) { return true; diff --git a/tensorflow/compiler/mlir/lite/transforms/passes.h b/tensorflow/compiler/mlir/lite/transforms/passes.h index c9e71b5b75d6c5..4c44edc6e4c844 100644 --- a/tensorflow/compiler/mlir/lite/transforms/passes.h +++ b/tensorflow/compiler/mlir/lite/transforms/passes.h @@ -123,6 +123,9 @@ std::unique_ptr> CreateQuantizePass( std::unique_ptr> CreateDefaultQuantizePass(); +std::unique_ptr> +CreateFoldStablehloConstantTransformsPass(); + std::unique_ptr> CreateLowerQuantAnnotationsPass(); // Creates an instance of the TFLite PropagateQParams pass which propagates diff --git a/tensorflow/compiler/mlir/lite/transforms/passes.td b/tensorflow/compiler/mlir/lite/transforms/passes.td index 5aab64ac05b4d9..a82604770e9ad5 100644 --- a/tensorflow/compiler/mlir/lite/transforms/passes.td +++ b/tensorflow/compiler/mlir/lite/transforms/passes.td @@ -354,6 +354,21 @@ def LowerQuantAnnotationsPass : Pass<"tfl-lower-quant-annotations", "mlir::Modul ]; } +def FoldStablehloConstantTransformsPass : Pass<"tfl-fold-stablehlo-constant-transforms", "mlir::func::FuncOp"> { + let summary = "Folds stablehlo.transpose and stablehlo.reshape operations on constants."; + let description = [{ + Unlike the upstream folders, this handles `DenseResourceElementsAttr` + operands, which is what large weights are stored as by the time they + reach this point in the pipeline. + }]; + let constructor = "CreateFoldStablehloConstantTransformsPass()"; + let dependentDialects = [ + "TFL::TensorFlowLiteDialect", + "mlir::arith::ArithDialect", + "stablehlo::StablehloDialect" + ]; +} + def PropagateQParamsPass : Pass<"tfl-propagate-qparams", "mlir::ModuleOp"> { let summary = "Propagates Quantization Parameters (scale and zero point) information through the graph."; let description = [{ diff --git a/tensorflow/core/common_runtime/process_util_test.cc b/tensorflow/core/common_runtime/process_util_test.cc index d1a53e8f411854..20a190c839967b 100644 --- a/tensorflow/core/common_runtime/process_util_test.cc +++ b/tensorflow/core/common_runtime/process_util_test.cc @@ -28,6 +28,12 @@ TEST(ProcessUtilTest, NumThreads) { EXPECT_EQ(10, NumInterOpThreadsFromSessionOptions(opts)); } +TEST(ProcessUtilTest, NumThreadsNegative) { + SessionOptions opts; + opts.config.set_inter_op_parallelism_threads(-1); + EXPECT_GT(NumInterOpThreadsFromSessionOptions(opts), 0); +} + TEST(ProcessUtilTest, ThreadPool) { SessionOptions opts; opts.config.set_inter_op_parallelism_threads(10); diff --git a/tensorflow/core/kernels/fractional_max_pool_op.cc b/tensorflow/core/kernels/fractional_max_pool_op.cc index 5ae827dc25032a..939debca0fd02c 100644 --- a/tensorflow/core/kernels/fractional_max_pool_op.cc +++ b/tensorflow/core/kernels/fractional_max_pool_op.cc @@ -344,6 +344,10 @@ class FractionalMaxPoolGradOp : public OpKernel { for (int64_t hs = 0; hs < height_seq_tensor.dim_size(0) - 1; ++hs) { // height start and end. const int64_t height_start = height_seq_tensor_flat(hs); + OP_REQUIRES(context, height_start >= 0, + absl::InvalidArgumentError(absl::StrCat( + "Row sequence tensor values must not be negative, got ", + height_seq_tensor.DebugString()))); int64_t height_end = overlapping_ ? height_seq_tensor_flat(hs + 1) : height_seq_tensor_flat(hs + 1) - 1; height_end = std::min(height_end, height_max); @@ -354,6 +358,11 @@ class FractionalMaxPoolGradOp : public OpKernel { (b * output_size[1] + hs) * output_size[2] + ws; // width start and end. const int64_t width_start = width_seq_tensor_flat(ws); + OP_REQUIRES( + context, width_start >= 0, + absl::InvalidArgumentError(absl::StrCat( + "Col sequence tensor values must not be negative, got ", + width_seq_tensor.DebugString()))); int64_t width_end = overlapping_ ? width_seq_tensor_flat(ws + 1) : width_seq_tensor_flat(ws + 1) - 1; width_end = std::min(width_end, width_max); diff --git a/tensorflow/core/kernels/lookup_table_op.cc b/tensorflow/core/kernels/lookup_table_op.cc index f882223f760e59..e4b0ae6dddf91c 100644 --- a/tensorflow/core/kernels/lookup_table_op.cc +++ b/tensorflow/core/kernels/lookup_table_op.cc @@ -1042,6 +1042,11 @@ class LookupTableExportOp : public LookupTableOpKernel { OP_REQUIRES_OK(ctx, GetTable(ctx, &table)); core::ScopedUnref unref_me(table); + DataTypeVector expected_inputs = {expected_input_0_}; + DataTypeVector expected_outputs = {table->key_dtype(), + table->value_dtype()}; + OP_REQUIRES_OK(ctx, ctx->MatchSignature(expected_inputs, expected_outputs)); + OP_REQUIRES_OK(ctx, table->ExportValues(ctx)); } }; diff --git a/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.h b/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.h index 2221347395abaa..5ae1a278387e77 100644 --- a/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.h +++ b/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.h @@ -22,6 +22,7 @@ limitations under the License. #include "xla/tsl/platform/statusor.h" #include "tensorflow/core/common_runtime/gpu/gpu_event_mgr.h" +#include "tensorflow/core/framework/bounds_check.h" #include "tensorflow/core/framework/register_types.h" #include "tensorflow/core/kernels/gpu_prim.h" #include "tensorflow/core/kernels/gpu_prim_helpers.h" @@ -987,6 +988,46 @@ absl::Status SparseSegmentReductionFunctor::operator()( /*output=*/output.data()); } +// Finds the position of an out-of-range `indices` value (against `noutput`) +// and/or `segment_ids` value (against `nsegments`), if any. `bad_index_pos` +// and `bad_segment_pos` must already be initialized to +// std::numeric_limits::max(); a value is left unchanged if no +// violation of that kind is found, or set to the position of one such +// violation (not necessarily the first in `i` order) otherwise. +template +__global__ void SparseSegmentGradBoundsCheckKernel( + Index nouter, Index noutput, SegmentId nsegments, + const Index* __restrict__ indices_vec, // [nouter] + const SegmentId* __restrict__ segment_vec, // [nouter] + int32_t* __restrict__ bad_index_pos, // [1] + int32_t* __restrict__ bad_segment_pos) { // [1] + for (Index i : GpuGridRangeX(nouter)) { + if (!FastBoundsCheck(indices_vec[i], noutput)) { + GpuAtomicMin(bad_index_pos, static_cast(i)); + } + if (!FastBoundsCheck(segment_vec[i], nsegments)) { + GpuAtomicMin(bad_segment_pos, static_cast(i)); + } + } +} + +template +absl::Status LaunchSparseSegmentGradBoundsCheckKernel( + const GPUDevice& d, Index nouter, Index noutput, SegmentId nsegments, + const Index* indices_vec, const SegmentId* segment_vec, + int32_t* bad_index_pos, int32_t* bad_segment_pos) { + TF_ASSIGN_OR_RETURN( + GpuLaunchConfig64 config, + GetGpuLaunchConfig64( + nouter, d, &SparseSegmentGradBoundsCheckKernel, + /*dynamic_shared_memory_size=*/0, + /*block_size_limit=*/0)); + return GpuLaunchKernel(SparseSegmentGradBoundsCheckKernel, + config.block_count, config.thread_per_block, 0, + d.stream(), nouter, noutput, nsegments, indices_vec, + segment_vec, bad_index_pos, bad_segment_pos); +} + template struct SparseSegmentGradFunctor { void operator()(OpKernelContext* context, @@ -1003,6 +1044,115 @@ struct SparseSegmentGradFunctor { const Index nouter = indices_vec.dimension(0); const Index noutput = output_flat.dimension(0); + // The CPU functor bounds-checks every `indices` value against `noutput` + // and every `segment_ids` value against `nsegments` before using them to + // read/write `input`/`output`, returning an InvalidArgumentError on + // failure. This GPU path performed no such validation, so an + // out-of-range indices or segment_ids value caused an out-of-bounds + // read or write on the device instead of a normal, recoverable error. + // Validate here, without copying the full `indices`/`segment_ids` + // arrays to the host: a device-side kernel does the O(nouter) scan, and + // only a couple of small flags come back over the (synchronous) D2H + // transfer in the common, valid-input case. + if (nouter > 0) { + // The bounds-check kernel below packs each violation's position into + // an int32 (matching SparseSegmentGradV2Functor's existing + // `nouter`/`ninner` handling above, which notes that neither is + // expected to be huge), so reject inputs where a position wouldn't + // fit rather than silently truncating it. + OP_REQUIRES( + context, + static_cast(nouter) <= std::numeric_limits::max(), + absl::InvalidArgumentError( + absl::StrCat("Indices vector of length ", nouter, + " is too large to fit in int32."))); + + se::Stream* stream = context->op_device_context()->stream(); + OP_REQUIRES(context, stream != nullptr, + absl::InternalError("No GPU stream available.")); + + // Fast check, matching the CPU functor: `segment_vec` is assumed + // sorted in non-decreasing order, so its last element is the maximum + // segment id. A single element is enough to validate `nsegments` + // against the well-formed/sorted case and produce the same + // "Invalid number of segments" message the CPU functor uses for it. + ScratchSpace last_segment_id_host(context, 1, + /*on_host=*/true); + OP_REQUIRES_OK( + context, + stream->Memcpy( + last_segment_id_host.mutable_data(), + stream_executor::DeviceAddressBase( + const_cast(segment_vec.data()) + (nouter - 1), + sizeof(SegmentId)), + sizeof(SegmentId))); + OP_REQUIRES_OK(context, stream->BlockHostUntilDone()); + OP_REQUIRES(context, *last_segment_id_host.data() + 1 <= nsegments, + absl::InvalidArgumentError("Invalid number of segments")); + + // Full per-element validation, matching the CPU functor: catches + // malformed input that isn't actually sorted (so the fast check above + // can't tell it's invalid from the last element alone). Runs as a + // device-side scan; `bad_positions` only ever carries 2 int32s back to + // the host, regardless of `nouter`. + Tensor bad_positions; + OP_REQUIRES_OK(context, context->allocate_temp(DT_INT32, TensorShape({2}), + &bad_positions)); + int32_t* const bad_positions_ptr = bad_positions.flat().data(); + int32_t* const bad_index_pos_ptr = bad_positions_ptr; + int32_t* const bad_segment_pos_ptr = bad_positions_ptr + 1; + stream_executor::DeviceAddressBase bad_positions_device( + bad_positions_ptr, 2 * sizeof(int32_t)); + OP_REQUIRES_OK(context, + stream->Memset32(&bad_positions_device, + std::numeric_limits::max(), + 2 * sizeof(int32_t))); + OP_REQUIRES_OK(context, LaunchSparseSegmentGradBoundsCheckKernel( + device, nouter, noutput, nsegments, + indices_vec.data(), segment_vec.data(), + bad_index_pos_ptr, bad_segment_pos_ptr)); + ScratchSpace bad_positions_host(context, 2, /*on_host=*/true); + OP_REQUIRES_OK(context, + stream->Memcpy(bad_positions_host.mutable_data(), + bad_positions_device, 2 * sizeof(int32_t))); + OP_REQUIRES_OK(context, stream->BlockHostUntilDone()); + const int32_t bad_index_pos = bad_positions_host.data()[0]; + const int32_t bad_segment_pos = bad_positions_host.data()[1]; + + if (bad_index_pos != std::numeric_limits::max()) { + ScratchSpace bad_index_host(context, 1, /*on_host=*/true); + OP_REQUIRES_OK( + context, stream->Memcpy(bad_index_host.mutable_data(), + stream_executor::DeviceAddressBase( + const_cast(indices_vec.data()) + + bad_index_pos, + sizeof(Index)), + sizeof(Index))); + OP_REQUIRES_OK(context, stream->BlockHostUntilDone()); + OP_REQUIRES(context, false, + absl::InvalidArgumentError( + absl::StrCat("Index ", *bad_index_host.data(), + " out of range [0, ", noutput, ")."))); + } + if (bad_segment_pos != std::numeric_limits::max()) { + ScratchSpace bad_segment_host(context, 1, + /*on_host=*/true); + OP_REQUIRES_OK( + context, + stream->Memcpy(bad_segment_host.mutable_data(), + stream_executor::DeviceAddressBase( + const_cast(segment_vec.data()) + + bad_segment_pos, + sizeof(SegmentId)), + sizeof(SegmentId))); + OP_REQUIRES_OK(context, stream->BlockHostUntilDone()); + OP_REQUIRES(context, false, + absl::InvalidArgumentError( + absl::StrCat("Segment id ", *bad_segment_host.data(), + " out of range [0, ", nsegments, ")."))); + } + } + // Allocate and compute segment weights (for Mean/SqrtN operations only). Tensor weights; using Tweights = typename RealTypeIfComplex::type; diff --git a/tensorflow/core/kernels/segment_reduction_ops_impl.h b/tensorflow/core/kernels/segment_reduction_ops_impl.h index 79c5e531786ee5..6039f589c63c70 100644 --- a/tensorflow/core/kernels/segment_reduction_ops_impl.h +++ b/tensorflow/core/kernels/segment_reduction_ops_impl.h @@ -519,12 +519,34 @@ class UnsortedSegmentReductionOp : public OpKernel { internal::ValidateUnsortedSegmentReduction( this, context, data, segment_ids, num_segments)); const auto segment_flat = segment_ids.flat(); - const Index output_rows = internal::SubtleMustCopy(static_cast( - num_segments.dtype() == DT_INT32 ? num_segments.scalar()() - : num_segments.scalar()())); - OP_REQUIRES(context, output_rows >= 0, - errors::InvalidArgument("Input num_segments == ", output_rows, - " must not be negative.")); + // Copy the value out of the (possibly shared) input buffer before any + // validation, so that the value that is checked is the same one that is + // subsequently used (see SubtleMustCopy). + const int64_t num_segments_val = + num_segments.dtype() == DT_INT32 + ? static_cast( + internal::SubtleMustCopy(num_segments.scalar()())) + : internal::SubtleMustCopy(num_segments.scalar()()); + OP_REQUIRES(context, num_segments_val >= 0, + absl::InvalidArgumentError( + absl::StrCat("Input num_segments == ", num_segments_val, + " must not be negative."))); + // Guard against excessively large num_segments that would cause integer + // overflow in output tensor size calculations, leading to illegal memory + // access or process abort (see #117549). + // + // Use kint32max for every index type. Although an int64 index could + // represent 2^31, that value is not representable as a signed int32 and + // would wrap if any downstream code performs 32-bit indexing; a uniform + // bound avoids that edge case. Nothing practical is lost, since the CPU + // functor already allocates std::vector row_counter(num_segments). + const int64_t kMaxSegments = + static_cast(std::numeric_limits::max()); + OP_REQUIRES(context, num_segments_val <= kMaxSegments, + absl::InvalidArgumentError(absl::StrCat( + "Input num_segments == ", num_segments_val, + " is too large. Must be at most ", kMaxSegments))); + const Index output_rows = static_cast(num_segments_val); TensorShape output_shape; OP_REQUIRES_OK(context, output_shape.AddDimWithStatus(output_rows)); for (int i = segment_ids.dims(); i < data.dims(); i++) { diff --git a/tensorflow/core/tfrt/ifrt/BUILD b/tensorflow/core/tfrt/ifrt/BUILD index 2dbf203aef7204..ead62b4e02a110 100644 --- a/tensorflow/core/tfrt/ifrt/BUILD +++ b/tensorflow/core/tfrt/ifrt/BUILD @@ -103,25 +103,34 @@ cc_library( ], ) -tf_cc_test( - name = "ifrt_device_utils_test", - srcs = [ - "ifrt_device_utils_test.cc", - ], - tags = ["no_oss"], - deps = [ - ":ifrt_device_utils", - "@com_google_absl//absl/status", - "@com_google_absl//absl/status:status_matchers", - "@com_google_googletest//:gtest_main", - "@xla//xla/python/ifrt", - "@xla//xla/python/ifrt:attribute_map", - "@xla//xla/python/ifrt:mock", - "@xla//xla/service:device_assignment", - "@xla//xla/tsl/platform:status_matchers", - "@xla//xla/tsl/platform:statusor", - ], -) +# copybara:uncomment_begin(PjRt-IFRT TPU client is not available in OSS) +# # TODO(siqiaowu): ifrt_device_utils_test couples with TPU assumptions ("coords" +# # and "core_on_chip" device attributes) and requires a TPU test client, which is +# # not available in OSS. +# tf_cc_test( +# name = "ifrt_device_utils_test", +# srcs = [ +# "ifrt_device_utils_test.cc", +# ], +# tags = [ +# "no_oss", +# "requires-tpu", +# ], +# deps = [ +# ":ifrt_device_utils", +# "//learning/brain/research/ifrt/pjrt_ifrt/pjrt_tpu:pjrt_tpu_client_test_lib", +# "@com_google_googletest//:gtest_main", +# "@com_google_absl//absl/status", +# "@com_google_absl//absl/status:status_matchers", +# "@xla//xla/python/ifrt", +# "@xla//xla/python/ifrt:attribute_map", +# "@xla//xla/python/ifrt:test_util", +# "@xla//xla/service:device_assignment", +# "@xla//xla/tsl/platform:status_matchers", +# "@xla//xla/tsl/platform:statusor", +# ], +# ) +# copybara:uncomment_end cc_library( name = "ifrt_serving_executable", diff --git a/tensorflow/core/tfrt/ifrt/ifrt_device_utils_test.cc b/tensorflow/core/tfrt/ifrt/ifrt_device_utils_test.cc index 49f62d74cdece6..871af37f23f540 100644 --- a/tensorflow/core/tfrt/ifrt/ifrt_device_utils_test.cc +++ b/tensorflow/core/tfrt/ifrt/ifrt_device_utils_test.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/core/tfrt/ifrt/ifrt_device_utils.h" +#include #include #include #include @@ -24,8 +25,9 @@ limitations under the License. #include "absl/status/status.h" #include "absl/status/status_matchers.h" #include "xla/python/ifrt/attribute_map.h" +#include "xla/python/ifrt/client.h" #include "xla/python/ifrt/device.h" -#include "xla/python/ifrt/mock.h" +#include "xla/python/ifrt/test_util.h" #include "xla/service/device_assignment.h" #include "xla/tsl/platform/status_matchers.h" #include "xla/tsl/platform/statusor.h" @@ -34,99 +36,93 @@ namespace tensorflow { namespace ifrt_serving { namespace { using ::testing::ElementsAre; -using ::testing::Return; -using ::testing::ReturnRef; static constexpr int kNumReplicas = 1; static constexpr int kNumCoresPerReplica = 2; // Intentionally have more devices than kNumReplicas * kNumCoresPerReplica for // testing purposes. -static constexpr int kNumDevices = 4; -static constexpr int kDeviceIdOffset = 8; +static constexpr int kNumDevices = 2; class IfrtDeviceUtilsTest : public ::testing::Test { protected: void SetUp() override { - mocked_devices_.reserve(kNumDevices); - devices_.reserve(kNumDevices); - for (int i = 0; i < kNumDevices; ++i) { - mocked_devices_.push_back(std::make_unique()); - ON_CALL(*mocked_devices_[i], Attributes()) - .WillByDefault(ReturnRef(device_attributes_maps_[i])); - ON_CALL(*mocked_devices_[i], Id()) - .WillByDefault(Return(xla::ifrt::DeviceId(kDeviceIdOffset + i))); - ON_CALL(client_, LookupDevice(xla::ifrt::DeviceId(kDeviceIdOffset + i))) - .WillByDefault(Return(mocked_devices_[i].get())); - - devices_.push_back(mocked_devices_[i].get()); - }; - - ON_CALL(client_, devices()).WillByDefault(Return(devices_)); - - // Default use the last two devices. - xla::DeviceAssignment assignment(kNumReplicas, kNumCoresPerReplica); - assignment(0, 0) = kDeviceIdOffset + 2; - assignment(0, 1) = kDeviceIdOffset + 3; - - ON_CALL(client_, - GetDefaultDeviceAssignment(kNumReplicas, kNumCoresPerReplica)) - .WillByDefault(Return(assignment)); + ASSERT_OK_AND_ASSIGN(client_, xla::ifrt::test_util::GetClient()); + for (auto* device : client_->devices()) { + ASSERT_OK_AND_ASSIGN( + auto coords, + device->Attributes().Get>("coords")); + ASSERT_OK_AND_ASSIGN(auto core, + device->Attributes().Get("core_on_chip")); + ASSERT_EQ(coords.size(), 3); + devices_.push_back(device); + device_coords_.push_back( + {static_cast(coords[0]), static_cast(coords[1]), + static_cast(coords[2]), static_cast(core)}); + } + ASSERT_GE(devices_.size(), kNumDevices); } - xla::ifrt::MockClient client_; - std::vector> mocked_devices_; - + std::shared_ptr client_; std::vector devices_; - std::vector device_attributes_maps_ = { - xla::ifrt::AttributeMap(xla::ifrt::AttributeMap::Map{ - {"coords", xla::ifrt::AttributeMap::Int64ListValue({1, 0, 0})}, - {"core_on_chip", xla::ifrt::AttributeMap::Int64Value(0)}}), - xla::ifrt::AttributeMap(xla::ifrt::AttributeMap::Map{ - {"coords", xla::ifrt::AttributeMap::Int64ListValue({1, 0, 0})}, - {"core_on_chip", xla::ifrt::AttributeMap::Int64Value(1)}}), - xla::ifrt::AttributeMap(xla::ifrt::AttributeMap::Map{ - {"coords", xla::ifrt::AttributeMap::Int64ListValue({2, 0, 0})}, - {"core_on_chip", xla::ifrt::AttributeMap::Int64Value(0)}}), - xla::ifrt::AttributeMap(xla::ifrt::AttributeMap::Map{ - {"coords", xla::ifrt::AttributeMap::Int64ListValue({2, 0, 0})}, - {"core_on_chip", xla::ifrt::AttributeMap::Int64Value(1)}}), - }; + std::vector> device_coords_; }; TEST_F(IfrtDeviceUtilsTest, Basic) { - std::vector device_assignment_attr = {1, 0, 0, 1, 1, 0, 0, 0}; + std::vector device_assignment_attr; + device_assignment_attr.insert(device_assignment_attr.end(), + device_coords_[1].begin(), + device_coords_[1].end()); + device_assignment_attr.insert(device_assignment_attr.end(), + device_coords_[0].begin(), + device_coords_[0].end()); TF_ASSERT_OK_AND_ASSIGN( auto devices_from_attribute, - GetAssignedIfrtDevices(client_, kNumReplicas, kNumCoresPerReplica, + GetAssignedIfrtDevices(*client_, kNumReplicas, kNumCoresPerReplica, device_assignment_attr)); EXPECT_THAT(devices_from_attribute, ElementsAre(devices_[1], devices_[0])); } -TEST_F(IfrtDeviceUtilsTest, SeparateXCoordinates) { - std::vector device_assignment_attr = {1, 0, 0, 1, 2, 0, 0, 0}; +TEST_F(IfrtDeviceUtilsTest, InvertCoordinates) { + std::vector device_assignment_attr; + device_assignment_attr.insert(device_assignment_attr.end(), + device_coords_[0].begin(), + device_coords_[0].end()); + device_assignment_attr.insert(device_assignment_attr.end(), + device_coords_[1].begin(), + device_coords_[1].end()); TF_ASSERT_OK_AND_ASSIGN( auto devices_from_attribute, - GetAssignedIfrtDevices(client_, kNumReplicas, kNumCoresPerReplica, + GetAssignedIfrtDevices(*client_, kNumReplicas, kNumCoresPerReplica, device_assignment_attr)); - EXPECT_THAT(devices_from_attribute, ElementsAre(devices_[1], devices_[2])); + EXPECT_THAT(devices_from_attribute, ElementsAre(devices_[0], devices_[1])); } TEST_F(IfrtDeviceUtilsTest, EmptyDeviceAssignmentShallReturnDefault) { + TF_ASSERT_OK_AND_ASSIGN( + xla::DeviceAssignment default_assignment, + client_->GetDefaultDeviceAssignment(kNumReplicas, kNumCoresPerReplica)); + TF_ASSERT_OK_AND_ASSIGN( + xla::ifrt::Device * dev0, + client_->LookupDevice(xla::ifrt::DeviceId(default_assignment(0, 0)))); + TF_ASSERT_OK_AND_ASSIGN( + xla::ifrt::Device * dev1, + client_->LookupDevice(xla::ifrt::DeviceId(default_assignment(0, 1)))); + TF_ASSERT_OK_AND_ASSIGN( auto devices_from_attribute, - GetAssignedIfrtDevices(client_, kNumReplicas, kNumCoresPerReplica, + GetAssignedIfrtDevices(*client_, kNumReplicas, kNumCoresPerReplica, std::nullopt)); - EXPECT_THAT(devices_from_attribute, ElementsAre(devices_[2], devices_[3])); + EXPECT_THAT(devices_from_attribute, ElementsAre(dev0, dev1)); } TEST_F(IfrtDeviceUtilsTest, MismatchCoordinatesShallFail) { - std::vector device_assignment_attr = {1, 0, 0, 1, 3, 0, 0, 0}; - auto status = GetAssignedIfrtDevices(client_, 1, 2, device_assignment_attr); + std::vector device_assignment_attr = {999, 999, 999, 999, + 999, 999, 999, 998}; + auto status = GetAssignedIfrtDevices(*client_, 1, 2, device_assignment_attr); EXPECT_THAT(status, absl_testing::StatusIs(absl::StatusCode::kFailedPrecondition)); } } // namespace - } // namespace ifrt_serving } // namespace tensorflow diff --git a/tensorflow/lite/core/api/flatbuffer_conversions.cc b/tensorflow/lite/core/api/flatbuffer_conversions.cc index f0d93ff9d2bfc6..ec33d526dc05b0 100644 --- a/tensorflow/lite/core/api/flatbuffer_conversions.cc +++ b/tensorflow/lite/core/api/flatbuffer_conversions.cc @@ -1590,6 +1590,11 @@ TfLiteStatus ParseFullyConnected(const Operator* op, TF_LITE_ENSURE_STATUS( ConvertTensorType(schema_params->quantized_bias_type(), ¶ms->quantized_bias_type, error_reporter)); + if (const auto* quant_spec = schema_params->quant_spec()) { + // Borrowed from the model buffer; see TfLiteFullyConnectedParams. + params->quant_spec = quant_spec->data(); + params->quant_spec_size = static_cast(quant_spec->size()); + } switch (schema_params->weights_format()) { case FullyConnectedOptionsWeightsFormat_DEFAULT: params->weights_format = kTfLiteFullyConnectedWeightsFormatDefault; diff --git a/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc b/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc index 75f78b09a0b108..9a5773b985ae43 100644 --- a/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc +++ b/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc @@ -2107,6 +2107,11 @@ class Subgraph { node_index); return kTfLiteError; } + if (params->quant_spec != nullptr && params->quant_spec_size > 0) { + TF_LITE_MAYBE_KERNEL_LOG(context, "unsupported quant_spec in node #%d", + node_index); + return kTfLiteError; + } return kTfLiteOk; } diff --git a/tensorflow/lite/examples/label_image/label_image_test.cc b/tensorflow/lite/examples/label_image/label_image_test.cc index b42769dc8eaefe..65a2e3a40da7d0 100644 --- a/tensorflow/lite/examples/label_image/label_image_test.cc +++ b/tensorflow/lite/examples/label_image/label_image_test.cc @@ -16,8 +16,10 @@ limitations under the License. #include "tensorflow/lite/examples/label_image/label_image.h" #include +#include #include #include +#include #include #include #include diff --git a/tensorflow/lite/tools/serialization/option_writer_generator.cc b/tensorflow/lite/tools/serialization/option_writer_generator.cc index 2ecc872174a202..04ddae4d16a289 100644 --- a/tensorflow/lite/tools/serialization/option_writer_generator.cc +++ b/tensorflow/lite/tools/serialization/option_writer_generator.cc @@ -468,6 +468,12 @@ void GenerateImportForOp(FILE* fp, const std::string& op_name, vector_name = "scatter_dims_to_operand_dims"; vector_size = "num_scatter_dims_to_operand_dims"; vector_type = "int64_t"; + } else if (elem_name == "quant_spec") { + // Opaque flexbuffers blob; see FullyConnectedOptions.quant_spec. + is_vector = true; + vector_name = "quant_spec"; + vector_size = "quant_spec_size"; + vector_type = "uint8_t"; } if (is_vector) { diff --git a/tensorflow/python/kernel_tests/data_structures/lookup_ops_test.py b/tensorflow/python/kernel_tests/data_structures/lookup_ops_test.py index 9aba606f69b6ff..701005520196aa 100644 --- a/tensorflow/python/kernel_tests/data_structures/lookup_ops_test.py +++ b/tensorflow/python/kernel_tests/data_structures/lookup_ops_test.py @@ -116,6 +116,25 @@ def testStaticHashTable(self, is_anonymous): ) self.assertCountEqual([0, 1, 2], self.evaluate(exported_values_tensor)) + def testExportSignatureMismatch(self, is_anonymous): + if is_anonymous and not tf2.enabled(): + self.skipTest(SKIP_ANONYMOUS_IN_TF1_REASON) + table = self.getHashTable()( + lookup_ops.KeyValueTensorInitializer( + constant_op.constant([0], dtype=dtypes.int64), + constant_op.constant([1], dtype=dtypes.int64), + ), + -1, + experimental_is_anonymous=is_anonymous, + ) + self.initialize_table(table) + with self.assertRaises((errors_impl.InvalidArgumentError, ValueError)): + self.evaluate( + gen_lookup_ops.lookup_table_export_v2( + table.resource_handle, Tkeys=dtypes.string, Tvalues=dtypes.string + ) + ) + def testStaticHashTableFindHighRank(self, is_anonymous): if is_anonymous and not tf2.enabled(): self.skipTest(SKIP_ANONYMOUS_IN_TF1_REASON) diff --git a/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py b/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py index fcf5ae14da605d..9a25788c1693ff 100644 --- a/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py +++ b/tensorflow/python/kernel_tests/math_ops/segment_reduction_ops_test.py @@ -660,7 +660,8 @@ def testBadNumSegments(self): unsorted = math_ops.unsorted_segment_sum( np.ones((3)), segment_ids=898042203, num_segments=num_segments) with self.assertRaisesOpError( - "Encountered overflow when multiplying | must not be negative" + "Encountered overflow when multiplying | must not be negative | " + "is too large" ): self.evaluate(unsorted) @@ -724,6 +725,67 @@ def testNaNPropagation(self): self.assertTrue(np.isnan(evaluated_max[0][0])) self.assertAllClose(evaluated_max[0][1], 3.0) + def testNumSegmentsOverflow(self): + """Regression test: an oversized num_segments raises InvalidArgumentError. + + Excessively large num_segments values (near INT64_MAX) cause integer + overflow in internal size calculations, leading to illegal memory access + or process abort. The kint32max upper bound prevents this. See #117549. + """ + data = constant_op.constant([1, 2, 3, 4], dtype=dtypes_lib.int32) + segment_ids = constant_op.constant([0, 1, 0, 1], dtype=dtypes_lib.int64) + # Just above the 2G limit (2^31 + 1). + num_segments = constant_op.constant((1 << 31) + 1, dtype=dtypes_lib.int64) + with self.assertRaisesRegex( + errors_impl.InvalidArgumentError, "is too large" + ): + self.evaluate( + math_ops.unsorted_segment_sum( + data=data, segment_ids=segment_ids, num_segments=num_segments + ) + ) + + def testNumSegmentsOverflow_Int32Index(self): + """num_segments == 2^31 must be rejected when the index type is int32. + + With a signed int32 index the largest representable segment count is + kint32max (2^31 - 1), so exactly 2^31 must be flagged as too large + rather than silently overflowing the cast to Index. See #117549. + """ + data = constant_op.constant([1, 2, 3, 4], dtype=dtypes_lib.int32) + segment_ids = constant_op.constant([0, 1, 0, 1], dtype=dtypes_lib.int32) + # Exactly the 2^31 boundary: one past what a signed int32 can represent. + num_segments = constant_op.constant(1 << 31, dtype=dtypes_lib.int64) + with self.assertRaisesRegex( + errors_impl.InvalidArgumentError, "is too large" + ): + self.evaluate( + math_ops.unsorted_segment_sum( + data=data, segment_ids=segment_ids, num_segments=num_segments + ) + ) + + def testNumSegmentsOverflow_Int64Index(self): + """num_segments == 2^31 must be rejected with an int64 index too. + + The limit is kint32max for every index type, not 2^31. The two differ by + one, and an earlier revision of this check allowed 2^31 whenever Index was + wide enough to hold it -- which left a value that is not representable in a + signed 32-bit integer reaching any downstream code that indexes with int32. + This pins the 64-bit path to the same bound. See #117549. + """ + data = constant_op.constant([1, 2, 3, 4], dtype=dtypes_lib.int32) + segment_ids = constant_op.constant([0, 1, 0, 1], dtype=dtypes_lib.int64) + num_segments = constant_op.constant(1 << 31, dtype=dtypes_lib.int64) + with self.assertRaisesRegex( + errors_impl.InvalidArgumentError, "is too large" + ): + self.evaluate( + math_ops.unsorted_segment_sum( + data=data, segment_ids=segment_ids, num_segments=num_segments + ) + ) + class SparseSegmentReductionHelper(SegmentReductionHelper): @@ -1302,7 +1364,7 @@ def testGradientIndicesInvalid1(self): ] segment_indices = [0, 1, 2, 2] tf_indices = [8, 3, 0, 10] - with self.session(use_gpu=False): + with self.session(): for tf_op in ops_list: s = tf_op(tf_x, tf_indices, segment_indices, 10) with self.assertRaisesOpError(r"Index 10 out of range \[0, 10\)"): @@ -1317,7 +1379,7 @@ def testGradientIndicesInvalid2(self): ] segment_indices = [0, 1, 2, 2] tf_indices = [8, 3, -1, 9] - with self.session(use_gpu=False): + with self.session(): for tf_op in ops_list: s = tf_op(tf_x, tf_indices, segment_indices, 10) with self.assertRaisesOpError(r"Index -1 out of range \[0, 10\)"): @@ -1333,7 +1395,7 @@ def testGradientSegmentsInvalid1(self): ] segment_indices = [0, 1, 1, 4] # 5 segments tf_indices = [8, 3, 0, 9] - with self.session(use_gpu=False): + with self.session(): for tf_op in ops_list: s = tf_op(tf_x, tf_indices, segment_indices, 10) with self.assertRaisesOpError("Invalid number of segments"): @@ -1348,7 +1410,7 @@ def testGradientSegmentsInvalid2(self): ] segment_indices = [0, 1, 2, 0] tf_indices = [8, 3, 0, 9] - with self.session(use_gpu=False): + with self.session(): for tf_op in ops_list: s = tf_op(tf_x, tf_indices, segment_indices, 10) with self.assertRaisesOpError(r"Segment id 1 out of range \[0, 1\)"): @@ -1363,7 +1425,7 @@ def testGradientSegmentsInvalid3(self): ] segment_indices = [-1, 0, 1, 1] tf_indices = [8, 3, 0, 9] - with self.session(use_gpu=False): + with self.session(): for tf_op in ops_list: s = tf_op(tf_x, tf_indices, segment_indices, 10) with self.assertRaisesOpError(r"Segment id -1 out of range \[0, 2\)"): diff --git a/third_party/xla/.github/workflows/clang_format.yml b/third_party/xla/.github/workflows/clang_format.yml index 359d5b78d2a2d7..e267d670b023d8 100644 --- a/third_party/xla/.github/workflows/clang_format.yml +++ b/third_party/xla/.github/workflows/clang_format.yml @@ -41,7 +41,7 @@ jobs: - name: "Fetch HEAD of main branch" run: git fetch origin main --depth=1 - name: Install uv dependencies - uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 + uses: astral-sh/setup-uv@bec219d24cd3e171d82865faccec33120bb574f4 # v10.1.0 - name: "Run clang-format check" env: TARGET_REF: "origin/${{ github.base_ref || 'main' }}" diff --git a/third_party/xla/.github/workflows/zizmor.yml b/third_party/xla/.github/workflows/zizmor.yml index 3bcae86a2acdb9..ee405f39080dfb 100644 --- a/third_party/xla/.github/workflows/zizmor.yml +++ b/third_party/xla/.github/workflows/zizmor.yml @@ -33,7 +33,7 @@ jobs: persist-credentials: false - name: Install uv dependencies - uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 + uses: astral-sh/setup-uv@bec219d24cd3e171d82865faccec33120bb574f4 # v10.1.0 - name: Run zizmor run: uvx zizmor --format=github .github xla diff --git a/third_party/xla/docs/contributing.md b/third_party/xla/docs/contributing.md index fa08305d0b2a5c..66a713ab9480bd 100644 --- a/third_party/xla/docs/contributing.md +++ b/third_party/xla/docs/contributing.md @@ -107,6 +107,41 @@ the XLA runner could load and execute the program. That means that there should be no pointers to `HloInstruction` or to other parts of the compiler or the `StreamExecutor`. +##### GPU Compatibility Window {#gpu-compatibility-window} + +XLA:GPU guarantees the following compatibility windows between the compiler +that produces a serialized program and the runtime that executes it: + +* **6 months backward compatibility**: a runtime must be able to execute + programs compiled by a version of XLA up to 6 months older. +* **2 weeks forward compatibility**: a program compiled by a newer version of + XLA must run on a runtime that is up to 2 weeks older. + +**Backward compatibility** means the runtime must keep executing everything +older compilers may have emitted, so existing thunks must not be removed from +the runtime. To retire a thunk, first stop emitting it from the compiler, and +only remove its runtime support 6 months later. + +**Forward compatibility** means the compiler must not emit anything the runtime +cannot yet deserialize or execute. Changes that affect the serialized thunk +sequence, such as adding a new thunk kind or a field the runtime must +understand, must therefore land in two stages: + +1. **Runtime first.** Add the thunk, its proto definitions, and + deserialization support, without emitting it from the compiler. +2. **Compiler after 2 weeks.** Once the runtime change has been submitted for + at least 2 weeks, enable emission in the compiler (e.g. thunk emitter, HLO + passes that produce the new op). + +There are two ways to do this: + +* Split the change into two PRs and send the compiler PR once the runtime + change is 2 weeks old. +* Land both in a single PR, but guard the compiler side behind a + [feature flag](#code-standards) that is off by default, and flip it in a + follow-up after 2 weeks. This also lets you test the full path end to end + in the meantime. + ### Code standards * *Coding style*: We follow [Google's code style guide](https://google.github.io/styleguide/). diff --git a/third_party/xla/docs/operation_semantics.md b/third_party/xla/docs/operation_semantics.md index 84f0a0cb5c177d..ffc81985ea8642 100644 --- a/third_party/xla/docs/operation_semantics.md +++ b/third_party/xla/docs/operation_semantics.md @@ -5170,6 +5170,49 @@ in detail on the [broadcasting page](broadcasting.md). For StableHLO information see [StableHLO - shift_right_logical](https://openxla.org/stablehlo/spec#shift_right_logical). +## Shuffle + +See also +[`XlaBuilder::Shuffle`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h). + +**`Shuffle(operand, dimensions, mode)`** + +Arguments | Type | Semantics +------------- | ------------------- | ------------------------------------- +`operand` | `XlaOp` | array of type T +`dimensions` | `ArraySlice` | dimensions to shuffle +`mode` | `ShuffleMode` | the shuffle pattern to apply, together with the attributes of that pattern + +Shuffles the elements of the `operand` array along the specified `dimensions`, +generating an output array of the same shape. Because a shuffle only moves +elements around, it never changes the shape of the operand. + +The `mode` selects which pattern is applied, and carries the attributes of that +pattern: + +Mode | Attributes | Semantics +-------- | ---------- | --------------------------------------------------- +`rotate` | `shifts` | Rotates the elements to the left by `shifts` + +### `rotate` + +Rotates the elements of the `operand` array along the specified `dimensions` to +the left by the corresponding `shifts`, generating an output array of the same +shape. Shifts outside of `[0, size(dimension))` are supported, and wrap around +cyclically. Each element of the output array at a multidimensional index is +retrieved from the operand array at a transformed index. The multidimensional +index is transformed by rotating the index along each specified dimension (i.e., +for a dimension of size N with a corresponding shift S, the element at output +index i is taken from operand index (i + S) % N). + +1-dimensional example: + +```cpp +let a = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}; +Shuffle(a, {0}, shuffle::Rotate({4})) +// Result: {4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 0, 1, 2, 3} +``` + ## Sign See also diff --git a/third_party/xla/third_party/shardy/temporary.patch b/third_party/xla/third_party/shardy/temporary.patch index 38c78a67e426b3..95a2a325ed502a 100644 --- a/third_party/xla/third_party/shardy/temporary.patch +++ b/third_party/xla/third_party/shardy/temporary.patch @@ -192,18 +192,3 @@ index 509398da..a9462de1 100644 ++ %fadd = fadd double %fmul26, %fmul29 ++ ret double %fadd ++} -diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl -index 08380e9e..d929da09 100644 ---- a/third_party/llvm/workspace.bzl -+++ b/third_party/llvm/workspace.bzl -@@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive") - - def repo(name): - """Imports LLVM.""" -- LLVM_COMMIT = "0038be81b8b5361112faf82e69d4590364ee81b9" -- LLVM_SHA256 = "6675de5b328d62a6bbbef606720ab110d32c55e4645d4fa503e1cfab15bfe878" -+ LLVM_COMMIT = "a22291abb68744cb0b70978eb506a7c5592a3309" -+ LLVM_SHA256 = "cf19b82f0dfbef1e62afe520ba4ee6de4f4823524d7804490e88d815f822c8e9" - - tf_http_archive( - name = name, diff --git a/third_party/xla/third_party/shardy/workspace.bzl b/third_party/xla/third_party/shardy/workspace.bzl index ce3f7fdc207531..dd328941030d9f 100644 --- a/third_party/xla/third_party/shardy/workspace.bzl +++ b/third_party/xla/third_party/shardy/workspace.bzl @@ -18,8 +18,8 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls") def repo(): - SHARDY_COMMIT = "7105c6a0e1d29141406a120b111e5525ff7b2e64" - SHARDY_SHA256 = "5f307f4e3acd7f7b7780bb9c7bf92a6cfe4108067f3976d80f43bce9c89c8f5a" + SHARDY_COMMIT = "08e079b1c5f5b0a3a0c696abc5c56b6ac436b835" + SHARDY_SHA256 = "b3de5417edb836ddcda6c1800bb0a5782c58110ff9567675a82ada9562d1033b" tf_http_archive( name = "shardy", diff --git a/third_party/xla/third_party/tensor_ir/patches/pow_op.patch b/third_party/xla/third_party/tensor_ir/patches/pow_op.patch new file mode 100644 index 00000000000000..d202fe1aa92cc3 --- /dev/null +++ b/third_party/xla/third_party/tensor_ir/patches/pow_op.patch @@ -0,0 +1,36 @@ +diff --git a/lib/Conversion/TensorToCudaTile/OutlinedComputeOps.cpp b/lib/Conversion/TensorToCudaTile/OutlinedComputeOps.cpp +--- a/lib/Conversion/TensorToCudaTile/OutlinedComputeOps.cpp ++++ b/lib/Conversion/TensorToCudaTile/OutlinedComputeOps.cpp +@@ -293,7 +293,7 @@ + getSignedness(exponentType), + cuda_tile::RoundingMode::NEAREST_EVEN); + } +- rewriter.replaceOpWithNewOp(op, operands[0], rhs); ++ rewriter.replaceOpWithNewOp(op, operands[0], rhs); + return success(); + }) + .Case([&](auto) { +diff --git a/test/Conversion/TensorToCudaTile/iterspace_matmul.mlir b/test/Conversion/TensorToCudaTile/iterspace_matmul.mlir +--- a/test/Conversion/TensorToCudaTile/iterspace_matmul.mlir ++++ b/test/Conversion/TensorToCudaTile/iterspace_matmul.mlir +@@ -267,7 +267,7 @@ + // CHECK: %[[RHS2:.*]] = reshape %[[RHS]] : [[TILE_RHS]] -> tile<2x128x16xf32> + // CHECK: %[[MMA:.*]] = mmaf %[[LHS2]], %[[RHS2]], %[[ACCUM]] + // CHECK: %[[RESHAPE:.*]] = reshape %[[MMA]] : tile<2x64x16xf32> -> [[TILE_RS:tile<1x1x2x8x8x4x4xf32>]] +-// CHECK: %[[TRANS:.*]] = permute %[[RESHAPE]] [2, 3, 0, 4, 5, 1, 6] : [[TILE_RS]] -> [[TILE_TR:tile<2x8x1x8x4x1x4xf32>]] ++// CHECK: %[[TRANS:.*]] = permute %[[RESHAPE]] [2, 3, {{[01]}}, 4, 5, {{[01]}}, 6] : [[TILE_RS]] -> [[TILE_TR:tile<2x8x1x8x4x1x4xf32>]] + // CHECK: %[[RESULT:.*]] = broadcast %[[TRANS]] : [[TILE_TR]] -> tile<2x8x1x8x4x2x4xf32> + // CHECK: %[[OUT_VIEW:.*]] = make_tensor_view %[[OUT]], shape = [8, 8, 8, 8, 4, 8, 4], strides = [65536, 8192, 1024, 128, 32, 4, 1] + // CHECK: %[[OUT_PVIEW:.*]] = make_partition_view %[[OUT_VIEW]] +diff --git a/test/Conversion/TensorToCudaTile/reduction_padding.mlir b/test/Conversion/TensorToCudaTile/reduction_padding.mlir +--- a/test/Conversion/TensorToCudaTile/reduction_padding.mlir ++++ b/test/Conversion/TensorToCudaTile/reduction_padding.mlir +@@ -238,7 +238,7 @@ + // CHECK: %[[BLOCK:.*]], %{{.*}}, %{{.*}} = get_tile_block_id : tile + // CHECK: %[[LOOP1:.*]] = for %[[IVAR1:.*]] in (%[[ZERO]] to %[[C2]], step %[[ONE]]) + // CHECK-SAME: iter_values(%[[IARG1:.*]] = %[[ACCUM]]) -> ([[TILE]]) +-// CHECK: %[[A_MUL:.*]] = muli %[[IVAR]], %[[C16]] : tile ++// CHECK: %[[A_MUL:.*]] = muli %[[IVAR1]], %[[C16]] : tile + // CHECK: %[[A_MUL1:.*]] = reshape %[[A_MUL]] : tile -> tile<1xi32> + // CHECK: %[[A_MUL2:.*]] = broadcast %[[A_MUL1]] : tile<1xi32> -> tile<16xi32> + // CHECK: %[[A_IOTA:.*]] = iota : tile<16xi32> diff --git a/third_party/xla/third_party/tensor_ir/patches/series b/third_party/xla/third_party/tensor_ir/patches/series index 9b81d4c8e0dfda..2e39ebeb534e1d 100644 --- a/third_party/xla/third_party/tensor_ir/patches/series +++ b/third_party/xla/third_party/tensor_ir/patches/series @@ -1,2 +1,3 @@ unused_variable.patch symbol_op_interface.patch +pow_op.patch diff --git a/third_party/xla/third_party/tensor_ir/tensor_ir.BUILD b/third_party/xla/third_party/tensor_ir/tensor_ir.BUILD index 3d591efd26df72..1a97f4faac5f5d 100644 --- a/third_party/xla/third_party/tensor_ir/tensor_ir.BUILD +++ b/third_party/xla/third_party/tensor_ir/tensor_ir.BUILD @@ -1,6 +1,22 @@ load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library") +load("@rules_cc//cc:cc_binary.bzl", "cc_binary") load("@rules_cc//cc:cc_library.bzl", "cc_library") +cc_binary( + name = "tensor_ir-tblgen", + srcs = glob([ + "tools/tensor_ir-tblgen/*.cpp", + "tools/tensor_ir-tblgen/*.h", + ]), + deps = [ + "@llvm-project//llvm:Support", + "@llvm-project//llvm:TableGen", + "@llvm-project//mlir:MlirTableGenMain", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:TableGen", + ], +) + exports_files(["LICENSE"]) td_library( @@ -14,10 +30,105 @@ td_library( "@llvm-project//mlir:GPUOpsTdFiles", "@llvm-project//mlir:InferTypeOpInterfaceTdFiles", "@llvm-project//mlir:OpBaseTdFiles", + "@llvm-project//mlir:PtrTdFiles", "@llvm-project//mlir:SideEffectInterfacesTdFiles", ], ) +gentbl_cc_library( + name = "TensorIRTypesIncGen", + tbl_outs = [ + ( + [ + "-gen-typedef-decls", + "-typedefs-dialect=nv_tensor_ir", + ], + "include/tensor_ir/Dialect/TensorTypes.h.inc", + ), + ( + [ + "-gen-typedef-defs", + "-typedefs-dialect=nv_tensor_ir", + ], + "include/tensor_ir/Dialect/TensorTypes.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/tensor_ir/Dialect/TensorTypes.td", + deps = [":TensorIRTdFiles"], +) + +td_library( + name = "TensorIROptionsTdFiles", + srcs = [ + "include/tensor_ir/Options/Enums.td", + "include/tensor_ir/Options/Options.td", + "include/tensor_ir/Options/OptionsBase.td", + ], + includes = ["include"], + deps = [ + ":TensorIRTdFiles", + "@llvm-project//mlir:OpBaseTdFiles", + ], +) + +gentbl_cc_library( + name = "TensorIROptionsEnumsIncGen", + tbl_outs = [ + ( + ["-gen-enum-decls"], + "include/tensor_ir/Options/OptionsEnums.h.inc", + ), + ( + ["-gen-enum-defs"], + "include/tensor_ir/Options/OptionsEnums.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/tensor_ir/Options/Enums.td", + deps = [":TensorIROptionsTdFiles"], +) + +gentbl_cc_library( + name = "TensorIRCompilerOptionsIncGen", + tbl_outs = [ + ( + ["-gen-component-options"], + "include/tensor_ir/Options/CompilerOptions.h.inc", + ), + ], + tblgen = ":tensor_ir-tblgen", + td_file = "include/tensor_ir/Options/Options.td", + deps = [":TensorIROptionsTdFiles"], +) + +cc_library( + name = "NVTensorIRCompilerOptions", + srcs = [ + "lib/Options/BytecodeVersionOptions.cpp", + "lib/Options/CompilerInvocation.cpp", + "lib/Options/Options.cpp", + ], + hdrs = [ + "include/tensor_ir/Options/BytecodeVersionOptions.h", + "include/tensor_ir/Options/Options.h", + "include/tensor_ir/Options/OptionsDetail.h", + "include/tensor_ir/Options/OptionsEnums.h", + ], + includes = ["include"], + visibility = ["//visibility:public"], + deps = [ + ":NVTensorIRDialect", + ":NVTensorIRSupport", + ":NVTensorIRUtils", + ":TensorIRCompilerOptionsIncGen", + ":TensorIROptionsEnumsIncGen", + "@cuda_tile//:CudaTileBytecode", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:Support", + ], +) + gentbl_cc_library( name = "TensorIRDialectIncGen", tbl_outs = [ @@ -204,19 +315,6 @@ cc_library( ], ) -cc_library( - name = "NVTensorIRCudaApi", - srcs = ["lib/Support/CudaApi.cpp"], - hdrs = ["include/tensor_ir/Support/CudaApi.h"], - includes = ["include"], - visibility = ["//visibility:public"], - deps = [ - ":NVTensorIRSupport", - "@llvm-project//llvm:Support", - "@local_config_cuda//cuda:cuda_headers", - ], -) - cc_library( name = "NVTensorIRDialect", srcs = [ @@ -224,6 +322,7 @@ cc_library( "lib/Dialect/TensorAttrs.cpp", "lib/Dialect/TensorDialect.cpp", "lib/Dialect/TensorOps.cpp", + "lib/Dialect/TensorTypes.cpp", ], hdrs = [ "include/tensor_ir/Dialect/TensorIR.h", @@ -240,14 +339,18 @@ cc_library( ":TensorIROpInterfacesIncGen", ":TensorIROpsCanonicalizationIncGen", ":TensorIROpsIncGen", + ":TensorIRTypesIncGen", "@llvm-project//llvm:Support", "@llvm-project//mlir:BytecodeOpInterface", "@llvm-project//mlir:ControlFlowInterfaces", + "@llvm-project//mlir:DialectUtils", "@llvm-project//mlir:FunctionInterfaces", "@llvm-project//mlir:GPUDialect", "@llvm-project//mlir:IR", "@llvm-project//mlir:InferTypeOpInterface", + "@llvm-project//mlir:PtrDialect", "@llvm-project//mlir:SideEffectInterfaces", + "@llvm-project//mlir:ViewLikeInterface", ], ) @@ -289,11 +392,14 @@ cc_library( includes = ["include"], visibility = ["//visibility:public"], deps = [ + ":NVTensorIRCompilerOptions", + ":NVTensorIRConversionOptions", ":NVTensorIRDialect", - ":NVTensorIRRuntime", + ":NVTensorIRKernelArgLayout", ":NVTensorIRUtils", "@llvm-project//llvm:Support", "@llvm-project//mlir:IR", + "@llvm-project//mlir:Support", ], ) @@ -305,176 +411,120 @@ cc_library( visibility = ["//visibility:public"], deps = [ ":NVTensorIRAnalysis", + ":NVTensorIRCompilerOptions", ":NVTensorIRDialect", ":NVTensorIRSupport", ":NVTensorIRUtils", ":TensorIRTransformPassesIncGen", "@llvm-project//llvm:Support", + "@llvm-project//mlir:ArithDialect", + "@llvm-project//mlir:DialectUtils", "@llvm-project//mlir:FuncDialect", + "@llvm-project//mlir:GPUDialect", "@llvm-project//mlir:IR", + "@llvm-project//mlir:MemRefDialect", "@llvm-project//mlir:Pass", + "@llvm-project//mlir:PtrDialect", + "@llvm-project//mlir:SCFDialect", + "@llvm-project//mlir:SideEffectInterfaces", "@llvm-project//mlir:Support", + "@llvm-project//mlir:TransformUtils", + "@llvm-project//mlir:Transforms", ], ) cc_library( - name = "NVTensorIRToCudaTileConversion", - srcs = glob(["lib/Conversion/TensorToCudaTile/*.cpp"]), - hdrs = [ - "include/tensor_ir/Conversion/TensorToCudaTile/Options.h", - "include/tensor_ir/Conversion/TensorToCudaTile/TensorToCudaTile.h", - "include/tensor_ir/Conversion/TensorToCudaTile/TensorToCudaTileInternal.h", - ], + name = "NVTensorIRConversionOptions", + hdrs = ["include/tensor_ir/Conversion/TensorToCudaTile/Options.h"], includes = ["include"], visibility = ["//visibility:public"], deps = [ - ":NVTensorIRAnalysis", - ":NVTensorIRDialect", - ":NVTensorIRSupport", - ":NVTensorIRUtils", - ":TensorToCudaTileConversionPassIncGen", - "@cuda_tile//:CudaTileDialect", + ":NVTensorIRCompilerOptions", "@llvm-project//llvm:Support", - "@llvm-project//mlir:DialectUtils", - "@llvm-project//mlir:IR", - "@llvm-project//mlir:Pass", - "@llvm-project//mlir:SideEffectInterfaces", - "@llvm-project//mlir:TransformUtils", ], ) cc_library( - name = "NVTensorIRCudaTilePipelines", - srcs = ["lib/Compiler/CudaTile/Pipelines.cpp"], - hdrs = ["include/tensor_ir/Compiler/CudaTile/Pipelines.h"], + name = "NVTensorIRKernelArgLayout", + hdrs = ["include/tensor_ir/Runtime/CudaTile/KernelArgLayout.h"], includes = ["include"], visibility = ["//visibility:public"], deps = [ - ":NVTensorIRDialect", - ":NVTensorIRToCudaTileConversion", - ":NVTensorIRTransform", - "@llvm-project//mlir:Pass", - "@llvm-project//mlir:Transforms", + ":NVTensorIRConversionOptions", + "@llvm-project//llvm:Support", ], ) cc_library( - name = "NVTensorIRRuntime", - srcs = ["lib/Runtime/CudaTileRuntimeKernel.cpp"], - hdrs = [ - "include/tensor_ir/Runtime/CudaTile/CudaTileRuntimeKernel.h", - "include/tensor_ir/Runtime/CudaTile/KernelArgLayout.h", - "include/tensor_ir/Runtime/CudaTile/KernelLaunchHelpers.h", - "include/tensor_ir/Runtime/CudaTile/RuntimeOperandAccessor.h", - "include/tensor_ir/Runtime/IRuntimeKernel.h", - "include/tensor_ir/Runtime/Types.h", - ], + name = "NVTensorIRRuntimeTypes", + hdrs = ["include/tensor_ir/Runtime/Types.h"], includes = ["include"], visibility = ["//visibility:public"], deps = [ - ":NVTensorIRCudaApi", ":NVTensorIRSupport", - "@cuda_tile//:CudaTileBytecode", - "@cuda_tile//:CudaTileDialect", - "@llvm-project//llvm:Support", - "@llvm-project//mlir:IR", - "@llvm-project//mlir:Support", - "@local_config_cuda//cuda:cuda_headers", ], ) cc_library( - name = "NVTensorIRReference", - srcs = glob(["lib/Reference/*.cpp"]), + name = "NVTensorIRToCudaTileConversion", + srcs = glob(["lib/Conversion/TensorToCudaTile/*.cpp"]), hdrs = [ - "include/tensor_ir/Reference/reference_graph.h", - "include/tensor_ir/Reference/reference_node.h", - "include/tensor_ir/Reference/simplified_tensor.h", - "include/tensor_ir/Reference/tensor_memory.h", - "lib/Reference/constant_utils.h", - ], - includes = [ - "include", - "lib/Reference", + "include/tensor_ir/Conversion/TensorToCudaTile/TensorToCudaTile.h", + "include/tensor_ir/Conversion/TensorToCudaTile/TensorToCudaTileInternal.h", ], + includes = ["include"], visibility = ["//visibility:public"], deps = [ - ":NVTensorIRCudaApi", + ":NVTensorIRAnalysis", + ":NVTensorIRCompilerOptions", + ":NVTensorIRConversionOptions", ":NVTensorIRDialect", - ":NVTensorIRRuntime", ":NVTensorIRSupport", ":NVTensorIRUtils", + ":TensorToCudaTileConversionPassIncGen", + "@cuda_tile//:CudaTileDialect", "@llvm-project//llvm:Support", - "@llvm-project//mlir:IR", - "@llvm-project//mlir:Support", - "@local_config_cuda//cuda:cuda_headers", - ], -) - -cc_library( - name = "NVTensorIRRegistration", - srcs = ["lib/Registration/Registration.cpp"], - hdrs = ["include/tensor_ir/Registration/Registration.h"], - includes = ["include"], - visibility = ["//visibility:public"], - deps = [ - ":NVTensorIRDialect", "@llvm-project//mlir:ArithDialect", + "@llvm-project//mlir:DialectUtils", "@llvm-project//mlir:FuncDialect", - "@llvm-project//mlir:FuncExtensions", + "@llvm-project//mlir:GPUDialect", "@llvm-project//mlir:IR", + "@llvm-project//mlir:MemRefDialect", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:PtrDialect", + "@llvm-project//mlir:SCFDialect", + "@llvm-project//mlir:SideEffectInterfaces", + "@llvm-project//mlir:TransformUtils", ], ) cc_library( - name = "NVTensorIRCompiler", - srcs = [ - "lib/Compiler/Compiler.cpp", - "lib/Compiler/CudaTile/CudaTileCompiler.cpp", - "lib/Compiler/CudaTile/CudaTileFrontend.cpp", - "lib/Compiler/CudaTile/TileIRAssembly.cpp", - ], - hdrs = [ - "include/tensor_ir/Compiler/CompileOptions.h", - "include/tensor_ir/Compiler/Compiler.h", - "include/tensor_ir/Compiler/CudaTile/CudaTileCompiler.h", - "include/tensor_ir/Compiler/CudaTile/CudaTileFrontend.h", - "include/tensor_ir/Compiler/CudaTile/TileIRAssembly.h", - ], + name = "NVTensorIRCudaTilePipelines", + srcs = ["lib/Compiler/CudaTile/Pipelines.cpp"], + hdrs = ["include/tensor_ir/Compiler/CudaTile/Pipelines.h"], includes = ["include"], visibility = ["//visibility:public"], deps = [ - ":NVTensorIRAnalysis", - ":NVTensorIRCudaTilePipelines", + ":NVTensorIRConversionOptions", ":NVTensorIRDialect", - ":NVTensorIRRegistration", - ":NVTensorIRRuntime", - ":NVTensorIRSupport", ":NVTensorIRToCudaTileConversion", - ":NVTensorIRUtils", - "@cuda_tile//:CudaTileBytecode", - "@cuda_tile//:CudaTileDialect", - "@llvm-project//llvm:Support", - "@llvm-project//mlir:ArithDialect", - "@llvm-project//mlir:IR", - "@llvm-project//mlir:Parser", + ":NVTensorIRTransform", + "@llvm-project//mlir:FuncDialect", "@llvm-project//mlir:Pass", - "@llvm-project//mlir:Support", + "@llvm-project//mlir:Transforms", ], ) cc_library( - name = "NVTensorIRCAPI", - srcs = ["lib/CAPI/TensorIR.cpp"], - hdrs = ["include/tensor_ir-c/TensorIR.h"], + name = "NVTensorIRTileIRAssembly", + srcs = ["lib/Compiler/CudaTile/TileIRAssembly.cpp"], + hdrs = ["include/tensor_ir/Compiler/CudaTile/TileIRAssembly.h"], includes = ["include"], visibility = ["//visibility:public"], deps = [ - ":NVTensorIRCompiler", - ":NVTensorIRDialect", - ":NVTensorIRRegistration", - ":NVTensorIRRuntime", + ":NVTensorIRSupport", + ":NVTensorIRUtils", + "@cuda_tile//:CudaTileBytecode", "@llvm-project//llvm:Support", - "@llvm-project//mlir:CAPIIRHeaders", ], ) diff --git a/third_party/xla/third_party/tensor_ir/workspace.bzl b/third_party/xla/third_party/tensor_ir/workspace.bzl index 4ed33e9d8047f6..f5bc5c190d1b80 100644 --- a/third_party/xla/third_party/tensor_ir/workspace.bzl +++ b/third_party/xla/third_party/tensor_ir/workspace.bzl @@ -5,8 +5,8 @@ load("//third_party:repo.bzl", "tf_mirror_urls") def repo(): """Imports Tensor IR.""" - TENSOR_IR_COMMIT = "63692d79629e6f32a1d8757695590a59e0adbafd" - TENSOR_IR_SHA256 = "b80794d7c2bfb1bc1ca432d892977becba8d35ec6c18c586acbd648ccc8074dd" + TENSOR_IR_COMMIT = "664d385996371db2a8de67b76f82fc89e2470109" + TENSOR_IR_SHA256 = "846a7bb37ef6a67b0224bd9cb25583ad0d84370a9d2565e3c5a8223c0dcf91f0" quilt_http_archive( name = "tensor_ir", diff --git a/third_party/xla/xla/BUILD b/third_party/xla/xla/BUILD index 05dc9d4612307c..b720b59902e79c 100644 --- a/third_party/xla/xla/BUILD +++ b/third_party/xla/xla/BUILD @@ -1192,6 +1192,29 @@ cc_library( ], ) +cc_library( + name = "shuffle", + srcs = ["shuffle.cc"], + hdrs = ["shuffle.h"], + visibility = ["//visibility:public"], + deps = [ + ":xla_data_proto_cc", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/types:span", + ], +) + +xla_cc_test( + name = "shuffle_test", + srcs = ["shuffle_test.cc"], + deps = [ + ":shuffle", + ":xla_data_proto_cc", + "//xla/tsl/platform:test_main", + "@com_google_googletest//:gtest", + ], +) + cc_library( name = "window_util", srcs = ["window_util.cc"], diff --git a/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.cc b/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.cc index b7c1aa55f3304f..fee83871c42993 100644 --- a/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.cc +++ b/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.cc @@ -407,6 +407,7 @@ void AddNewXtileToVectorPasses(mlir::OpPassManager& pm) { pm.addPass(xtile::createExpandXtileComplexOpsPass()); pm.addPass(xtile::createStablehloLowerToArithPass()); pm.addPass(xtile::createLegalizeUnsignedIntegersAsSignlessPass()); + pm.addPass(cpu::createLegalizeNarrowFloatStoragePass()); pm.addPass(mlir::createCanonicalizerPass()); pm.addPass(cpu::createVectorizeXTilePass()); pm.addPass(cpu::createLowerXTileEntryPass()); @@ -484,6 +485,8 @@ void AddNewVectorToLLVMPasses(mlir::OpPassManager& pm, bool fast_min_max, options.vectorTransposeLowering = mlir::vector::VectorTransposeLowering::Shuffle16x16; pm.addPass(mlir::createConvertVectorToLLVMPass(options)); + pm.addPass(cpu::createLowerMemRefBitcastPass()); + pm.addPass(cpu::createLowerToLLVMPass()); pm.addPass(mlir::memref::createExpandStridedMetadataPass()); pm.addPass(emitters::createSafeIntegerArithmeticPass()); diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/add_bf16_mask.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/add_bf16_mask.hlo new file mode 100644 index 00000000000000..412198446a0d55 --- /dev/null +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/add_bf16_mask.hlo @@ -0,0 +1,72 @@ +// RUN: fusion_to_xtile %s | fusion_compiler_opt --xtile-cpu-xtile-to-vector \ +// RUN: --xtile-cpu-vector-to-llvm | FileCheck %s --check-prefix=CHECK-LLVM +// RUN: test_correctness %s + +// RUN: fusion_to_xtile %s | fusion_compiler_opt \ +// RUN: --xtile-cpu-new-xtile-to-vector \ +// RUN: | FileCheck %s --check-prefix=NEW-VECTOR +// RUN: fusion_to_xtile %s | fusion_compiler_opt \ +// RUN: --xtile-cpu-new-xtile-to-vector --xtile-cpu-new-vector-to-llvm \ +// RUN: | FileCheck %s --check-prefix=NEW-VECTOR-LLVM +// RUN: test_correctness %s --xla_cpu_use_new_xtile_lowering=true + +// The tile does not divide the minor dimension, so the reads are guarded by an +// `scf.if` that picks between an in-bounds read and a masked read. None of the +// data movement may happen on bf16 vectors: x86 has no legal vector register +// class for bf16 without AVX512-BF16, so SelectionDAG would scalarize the merge +// and the masked accesses and emit one `__truncsfbf2` libcall per element. +// `xtile-cpu-legalize-narrow-float-storage` views the buffers as i16, so bf16 +// only shows up right at `arith.extf`/`arith.truncf`. + +fusion { + p0 = bf16[100,100] parameter(0) + p1 = bf16[100,100] parameter(1) + c0 = f32[100,100] convert(p0) + c1 = f32[100,100] convert(p1) + a = f32[100,100] add(c0, c1) + ROOT r = bf16[100,100] convert(a) +} + +// NEW-VECTOR-LABEL: func.func @wrapped_fusion_impl +// NEW-VECTOR-SAME: %[[ARG0:.*]]: memref<100x100xbf16>, %[[ARG1:.*]]: memref<100x100xbf16>, %[[ARG2:.*]]: memref<100x100xbf16> + +// NEW-VECTOR-DAG: %[[PAD:.*]] = arith.constant 0 : i16 +// NEW-VECTOR-DAG: %[[OUT:.*]] = xtile.memref_bitcast %[[ARG2]] : memref<100x100xbf16> -> memref<100x100xi16> +// NEW-VECTOR-DAG: %[[IN0:.*]] = xtile.memref_bitcast %[[ARG0]] : memref<100x100xbf16> -> memref<100x100xi16> +// NEW-VECTOR: scf.for +// NEW-VECTOR: %[[COND:.*]] = arith.cmpi sge +// NEW-VECTOR: %[[R0:.*]] = scf.if %[[COND]] -> (vector<1x16xi16>) { +// NEW-VECTOR: vector.transfer_read %[[IN0]]{{.*}}, %[[PAD]] {in_bounds = [true, true]} : memref<100x100xi16>, vector<1x16xi16> +// NEW-VECTOR: } else { +// NEW-VECTOR: %[[MASK:.*]] = vector.create_mask {{.*}} : vector<1x16xi1> +// NEW-VECTOR: vector.transfer_read %[[IN0]]{{.*}}, %[[PAD]], %[[MASK]] : memref<100x100xi16>, vector<1x16xi16> +// NEW-VECTOR: } +// NEW-VECTOR: %[[R0_BF16:.*]] = arith.bitcast %[[R0]] : vector<1x16xi16> to vector<1x16xbf16> +// NEW-VECTOR: arith.extf %[[R0_BF16]] : vector<1x16xbf16> to vector<1x16xf32> +// NEW-VECTOR: %[[TRUNC:.*]] = arith.truncf %{{.*}} : vector<1x16xf32> to vector<1x16xbf16> +// NEW-VECTOR: %[[TRUNC_I16:.*]] = arith.bitcast %[[TRUNC]] : vector<1x16xbf16> to vector<1x16xi16> +// NEW-VECTOR: vector.transfer_write %[[TRUNC_I16]], %[[OUT]]{{.*}} : vector<1x16xi16>, memref<100x100xi16> + +// NEW-VECTOR-LLVM-LABEL: llvm.func internal @wrapped_fusion_impl +// NEW-VECTOR-LLVM-NOT: xtile.memref_bitcast +// NEW-VECTOR-LLVM: ^bb{{[0-9]+}}(%{{.*}}: !llvm.array<1 x vector<16xi16>>) +// NEW-VECTOR-LLVM: ^bb{{[0-9]+}}(%{{.*}}: !llvm.array<1 x vector<16xi16>>) +// NEW-VECTOR-LLVM-NOT: llvm.intr.masked.{{.*}}vector<16xbf16> +// NEW-VECTOR-LLVM: llvm.intr.masked.store{{.*}}vector<16xi16> +// NEW-VECTOR-LLVM-NOT: llvm.intr.masked.{{.*}}vector<16xbf16> + +// CHECK-LLVM-LABEL: llvm.func internal @wrapped_fusion_impl + +ENTRY main { + p0 = bf16[100,100] parameter(0) + p1 = bf16[100,100] parameter(1) + ROOT fusion = bf16[100,100] fusion(p0, p1), kind=kLoop, calls=fusion, + backend_config={ + "fusion_backend_config":{ + "block_level_fusion_config":{ + "output_tiles":[{"sizes":["1", "16"]}], + "num_ctas":"1" + } + } + } +} diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/convert_f32_bf16.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/convert_f32_bf16.hlo index 85cd8460d6f1e6..0556729980ea28 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/convert_f32_bf16.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/convert_f32_bf16.hlo @@ -16,7 +16,7 @@ fusion { // NEW-VECTOR-LLVM-LABEL: llvm.func internal @wrapped_fusion_impl // NEW-VECTOR-LLVM: %[[LOAD:.*]] = llvm.load %{{.*}} : !llvm.ptr -> vector<16xf32> // NEW-VECTOR-LLVM-COUNT-2: llvm.call @xla.fptrunc.v8f32.to.v8bf16(%{{.*}}) : (vector<8xf32>) -> vector<8xbf16> -// NEW-VECTOR-LLVM: llvm.store %{{.*}}, %{{.*}} : vector<16xbf16>, !llvm.ptr +// NEW-VECTOR-LLVM: llvm.store %{{.*}}, %{{.*}} : vector<16xi16>, !llvm.ptr ENTRY main { p0 = f32[1, 16] parameter(0) diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/is_finite.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/is_finite.hlo index cdb8e7604f65f0..05965d7a9060e5 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/is_finite.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/is_finite.hlo @@ -42,9 +42,11 @@ // NEW-VECTOR-LABEL: func.func @wrapped_fusion_impl // NEW-VECTOR-SAME: %[[ARG:[^:]+]]: memref<1x9xf16>, %[[OUT:[^:]+]]: memref<1x9xi8>, %[[PID:[^:]+]]: index +// NEW-VECTOR: %[[ARG_I16:.*]] = xtile.memref_bitcast %[[ARG]] : memref<1x9xf16> -> memref<1x9xi16> // NEW-VECTOR: %[[MASK:.*]] = vector.constant_mask [1, 9] : vector<1x16xi1> -// NEW-VECTOR: %[[READ:.*]] = vector.transfer_read %[[ARG]]{{.*}}, %[[MASK]]{{.*}} : memref<1x9xf16>, vector<1x16xf16> -// NEW-VECTOR: %[[FINITE:.*]] = math.isfinite %[[READ]] : vector<1x16xf16> +// NEW-VECTOR: %[[READ:.*]] = vector.transfer_read %[[ARG_I16]]{{.*}}, %[[MASK]]{{.*}} : memref<1x9xi16>, vector<1x16xi16> +// NEW-VECTOR: %[[READ_F16:.*]] = arith.bitcast %[[READ]] : vector<1x16xi16> to vector<1x16xf16> +// NEW-VECTOR: %[[FINITE:.*]] = math.isfinite %[[READ_F16]] : vector<1x16xf16> // NEW-VECTOR: %[[CAST:.*]] = arith.extui %[[FINITE]] : vector<1x16xi1> to vector<1x16xi8> // NEW-VECTOR: vector.transfer_write %[[CAST]], %[[OUT]]{{.*}}, %[[MASK]]{{.*}} : vector<1x16xi8>, memref<1x9xi8> diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_bf16_f32.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_bf16_f32.hlo index 2ac596a5d7e59e..82c69b890bd617 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_bf16_f32.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_bf16_f32.hlo @@ -61,16 +61,18 @@ fusion { // CHECK-LLVM: llvm.alloca // NEW-VECTOR-LABEL: func.func @wrapped_fusion_impl -// NEW-VECTOR: %[[PAD:.*]] = arith.constant 0.000000e+00 : bf16 +// NEW-VECTOR: %[[PAD:.*]] = arith.constant 0 : i16 +// NEW-VECTOR: %[[IN:.*]] = xtile.memref_bitcast %{{.*}} : memref<1x256x56x56xbf16> -> memref<1x256x56x56xi16> // NEW-VECTOR: scf.for -// NEW-VECTOR: %[[READ:.*]] = vector.transfer_read %{{.*}}[%{{.*}}], %[[PAD]] {in_bounds = [true, true, true, true]} : memref<1x256x56x56xbf16>, vector<1x8x1x8xbf16> -// NEW-VECTOR: %[[EXT:.*]] = arith.extf %[[READ]] : vector<1x8x1x8xbf16> to vector<1x8x1x8xf32> +// NEW-VECTOR: %[[READ:.*]] = vector.transfer_read %[[IN]][%{{.*}}], %[[PAD]] {in_bounds = [true, true, true, true]} : memref<1x256x56x56xi16>, vector<1x8x1x8xi16> +// NEW-VECTOR: %[[CAST:.*]] = arith.bitcast %[[READ]] : vector<1x8x1x8xi16> to vector<1x8x1x8xbf16> +// NEW-VECTOR: %[[EXT:.*]] = arith.extf %[[CAST]] : vector<1x8x1x8xbf16> to vector<1x8x1x8xf32> // NEW-VECTOR: %[[TRANS:.*]] = vector.transpose %[[EXT]], [0, 2, 3, 1] : vector<1x8x1x8xf32> to vector<1x1x8x8xf32> // NEW-VECTOR: vector.transfer_write %[[TRANS]], %{{.*}}[%{{.*}}] {in_bounds = [true, true, true, true]} : vector<1x1x8x8xf32>, memref<1x56x56x256xf32> // NEW-VECTOR-LLVM-LABEL: llvm.func internal @wrapped_fusion_impl // NEW-VECTOR-LLVM-NOT: llvm.alloca -// NEW-VECTOR-LLVM-COUNT-8: llvm.load %{{.*}} : !llvm.ptr -> vector<8xbf16> +// NEW-VECTOR-LLVM-COUNT-8: llvm.load %{{.*}} : !llvm.ptr -> vector<8xi16> // NEW-VECTOR-LLVM: llvm.shufflevector // NEW-VECTOR-LLVM-COUNT-8: llvm.store %{{.*}}, %{{.*}} : vector<8xf32>, !llvm.ptr diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_f32_bf16.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_f32_bf16.hlo index c28222377502bf..cca9185f4e72bc 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_f32_bf16.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_f32_bf16.hlo @@ -65,15 +65,17 @@ fusion { // NEW-VECTOR: scf.for // NEW-VECTOR: %[[READ:.*]] = vector.transfer_read %{{.*}}[%{{.*}}], %[[PAD]] {in_bounds = [true, true, true, true]} : memref<1x56x56x256xf32>, vector<1x1x8x8xf32> // NEW-VECTOR: %[[TRUNC:.*]] = arith.truncf %[[READ]] : vector<1x1x8x8xf32> to vector<1x1x8x8xbf16> -// NEW-VECTOR: %[[TRANS:.*]] = vector.transpose %[[TRUNC]], [0, 3, 1, 2] : vector<1x1x8x8xbf16> to vector<1x8x1x8xbf16> -// NEW-VECTOR: vector.transfer_write %[[TRANS]], %{{.*}}[%{{.*}}] {in_bounds = [true, true, true, true]} : vector<1x8x1x8xbf16>, memref<1x256x56x56xbf16> +// NEW-VECTOR: %[[CAST:.*]] = arith.bitcast %[[TRUNC]] : vector<1x1x8x8xbf16> to vector<1x1x8x8xi16> +// NEW-VECTOR: %[[TRANS:.*]] = vector.transpose %[[CAST]], [0, 3, 1, 2] : vector<1x1x8x8xi16> to vector<1x8x1x8xi16> +// NEW-VECTOR: vector.transfer_write %[[TRANS]], %{{.*}}[%{{.*}}] {in_bounds = [true, true, true, true]} : vector<1x8x1x8xi16>, memref<1x256x56x56xi16> // NEW-VECTOR-LLVM-LABEL: llvm.func internal @wrapped_fusion_impl // NEW-VECTOR-LLVM-NOT: llvm.alloca // NEW-VECTOR-LLVM-COUNT-8: llvm.load %{{.*}} : !llvm.ptr -> vector<8xf32> // NEW-VECTOR-LLVM-COUNT-8: llvm.call @xla.fptrunc.v8f32.to.v8bf16 -// NEW-VECTOR-LLVM: llvm.shufflevector -// NEW-VECTOR-LLVM-COUNT-8: llvm.store %{{.*}}, %{{.*}} : vector<8xbf16>, !llvm.ptr +// NEW-VECTOR-LLVM-COUNT-8: llvm.bitcast %{{.*}} : vector<8xbf16> to vector<8xi16> +// NEW-VECTOR-LLVM: llvm.shufflevector %{{.*}}, %{{.*}} [{{.*}}] : vector<8xi16> +// NEW-VECTOR-LLVM-COUNT-8: llvm.store %{{.*}}, %{{.*}} : vector<8xi16>, !llvm.ptr ENTRY %main { %input = f32[1,56,56,256]{3,2,1,0} parameter(0) diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_broadcast_f16.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_broadcast_f16.hlo new file mode 100644 index 00000000000000..e32862353cdc8d --- /dev/null +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_broadcast_f16.hlo @@ -0,0 +1,40 @@ +// RUN: fusion_to_xtile %s | fusion_compiler_opt \ +// RUN: --xtile-cpu-new-xtile-to-vector \ +// RUN: | FileCheck %s --check-prefix=NEW-VECTOR +// RUN: test_correctness %s --xla_cpu_use_new_xtile_lowering=true + +// f16 has the same problem as bf16 on x86 without AVX512-FP16. The broadcast +// and the transpose only move data, so they run on i16; f16 is only visible +// at the compute. + +fusion { + p0 = f16[64,48] parameter(0) + p1 = f16[48] parameter(1) + t = f16[48,64] transpose(p0), dimensions={1,0} + b = f16[48,64] broadcast(p1), dimensions={0} + ROOT a = f16[48,64] add(t, b) +} + +// NEW-VECTOR-LABEL: func.func @wrapped_fusion_impl +// NEW-VECTOR-DAG: xtile.memref_bitcast %{{.*}} : memref<64x48xf16> -> memref<64x48xi16> +// NEW-VECTOR-DAG: xtile.memref_bitcast %{{.*}} : memref<48xf16> -> memref<48xi16> +// NEW-VECTOR-DAG: xtile.memref_bitcast %{{.*}} : memref<48x64xf16> -> memref<48x64xi16> +// NEW-VECTOR: vector.transpose %{{.*}} : vector<8x8xi16> to vector<8x8xi16> +// NEW-VECTOR: vector.broadcast %{{.*}} : vector<8x1xi16> to vector<8x8xi16> +// NEW-VECTOR-NOT: vector.transpose {{.*}}f16> +// NEW-VECTOR-NOT: vector.broadcast {{.*}}f16> +// NEW-VECTOR-NOT: vector.transfer_{{.*}}f16> + +ENTRY main { + p0 = f16[64,48] parameter(0) + p1 = f16[48] parameter(1) + ROOT fusion = f16[48,64] fusion(p0, p1), kind=kLoop, calls=fusion, + backend_config={ + "fusion_backend_config":{ + "block_level_fusion_config":{ + "output_tiles":[{"sizes":["8", "8"]}], + "num_ctas":"1" + } + } + } +} diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_convert_bf16.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_convert_bf16.hlo new file mode 100644 index 00000000000000..f69f4d8820ff57 --- /dev/null +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_convert_bf16.hlo @@ -0,0 +1,50 @@ +// RUN: fusion_to_xtile %s | fusion_compiler_opt \ +// RUN: --xtile-cpu-new-xtile-to-vector \ +// RUN: | FileCheck %s --check-prefix=NEW-VECTOR +// RUN: fusion_to_xtile %s | fusion_compiler_opt \ +// RUN: --xtile-cpu-new-xtile-to-vector --xtile-cpu-new-vector-to-llvm \ +// RUN: | FileCheck %s --check-prefix=NEW-VECTOR-LLVM +// RUN: test_correctness %s --xla_cpu_use_new_xtile_lowering=true + +// The output side of an NHWC -> NCHW layout change in bf16, as emitted around a +// depthwise convolution. The convert happens before the transpose, so without +// help the 16x16 transpose shuffles bf16 vectors, which x86 scalarizes into one +// `__truncsfbf2` libcall per element. The transpose only moves data, so it must +// run on i16 instead. + +fusion { + p0 = f32[1,56,56,256]{3,2,1,0} parameter(0) + c = bf16[1,56,56,256]{3,2,1,0} convert(p0) + t = bf16[1,256,56,56]{1,3,2,0} transpose(c), dimensions={0,3,1,2} + ROOT copy = bf16[1,256,56,56]{3,2,1,0} copy(t) +} + +// NEW-VECTOR-LABEL: func.func @wrapped_fusion_impl +// NEW-VECTOR: %[[OUT:.*]] = xtile.memref_bitcast %{{.*}} : memref<1x256x56x56xbf16> -> memref<1x256x56x56xi16> +// NEW-VECTOR: %[[TRUNC:.*]] = arith.truncf %{{.*}} : vector<1x1x16x16xf32> to vector<1x1x16x16xbf16> +// NEW-VECTOR: %[[CAST:.*]] = arith.bitcast %[[TRUNC]] : vector<1x1x16x16xbf16> to vector<1x1x16x16xi16> +// NEW-VECTOR: %[[T:.*]] = vector.transpose %[[CAST]], [0, 3, 1, 2] : vector<1x1x16x16xi16> to vector<1x16x1x16xi16> +// NEW-VECTOR: vector.transfer_write %[[T]], %[[OUT]] + +// The transpose shuffles and all memory accesses are on i16. The only bf16 +// vectors left join the two halves produced by the 8-wide truncation intrinsic +// right before the bitcast, which LLVM keeps in registers. +// NEW-VECTOR-LLVM-LABEL: llvm.func internal @wrapped_fusion_impl +// NEW-VECTOR-LLVM-NOT: llvm.shufflevector {{.*}}[0, 16, 1, 17, {{.*}}bf16> +// NEW-VECTOR-LLVM: llvm.shufflevector {{.*}}[0, 16, 1, 17, {{.*}} : vector<16xi16> +// NEW-VECTOR-LLVM-NOT: llvm.shufflevector {{.*}}[0, 16, 1, 17, {{.*}}bf16> +// NEW-VECTOR-LLVM-NOT: llvm.store {{.*}}bf16> +// NEW-VECTOR-LLVM-NOT: llvm.intr.masked.store {{.*}}bf16> + +ENTRY main { + p0 = f32[1,56,56,256]{3,2,1,0} parameter(0) + ROOT fusion = bf16[1,256,56,56]{3,2,1,0} fusion(p0), kind=kLoop, calls=fusion, + backend_config={ + "fusion_backend_config":{ + "block_level_fusion_config":{ + "output_tiles":[{"sizes":["1", "16", "1", "16"]}], + "num_ctas":"12" + } + } + } +} diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/BUILD b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/BUILD index 6dc1b55ac96e50..ce7c4c57144e07 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/BUILD +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/BUILD @@ -68,7 +68,9 @@ cc_library( "fuse_elementwise_pass.cc", "hoist_alloca.cc", "initialize_allocs.cc", + "legalize_narrow_float_storage.cc", "linalg_elementwise_to_vector_pass.cc", + "lower_memref_bitcast.cc", "lower_xtile_entry.cc", "memref_copy_to_loops.cc", "shlo_to_vector.cc", @@ -108,6 +110,7 @@ cc_library( "@llvm-project//mlir:FuncTransforms", "@llvm-project//mlir:FuncUtils", "@llvm-project//mlir:IR", + "@llvm-project//mlir:LLVMCommonConversion", "@llvm-project//mlir:LLVMDialect", "@llvm-project//mlir:LinalgDialect", "@llvm-project//mlir:LinalgTransforms", diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/legalize_narrow_float_storage.cc b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/legalize_narrow_float_storage.cc new file mode 100644 index 00000000000000..12a682cdcc04f8 --- /dev/null +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/legalize_narrow_float_storage.cc @@ -0,0 +1,228 @@ +/* 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 "llvm/ADT/SmallVector.h" +#include "llvm/ADT/SmallVectorExtras.h" +#include "llvm/Support/LogicalResult.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinOps.h" // IWYU pragma: keep +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/Operation.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/IR/Types.h" +#include "mlir/IR/Value.h" +#include "mlir/Pass/Pass.h" // IWYU pragma: keep +#include "mlir/Support/LLVM.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "stablehlo/dialect/StablehloOps.h" +#include "xla/backends/cpu/codegen/tiled/transforms/passes.h" // IWYU pragma: keep +#include "xla/codegen/xtile/ir/xtile_ops.h" + +namespace xla::cpu { + +#define GEN_PASS_DEF_LEGALIZENARROWFLOATSTORAGEPASS +#include "xla/backends/cpu/codegen/tiled/transforms/passes.h.inc" + +namespace { + +namespace ma = ::mlir::arith; +namespace shlo = ::mlir::stablehlo; +namespace tensor = ::mlir::tensor; + +using ::mlir::IntegerType; +using ::mlir::Location; +using ::mlir::LogicalResult; +using ::mlir::MemRefType; +using ::mlir::OpBuilder; +using ::mlir::PatternRewriter; +using ::mlir::RankedTensorType; +using ::mlir::Type; +using ::mlir::Value; + +// Returns the integer type used to carry values of `type` through data +// movement, or nullptr if `type` is not a narrow float. +IntegerType GetStorageType(Type type) { + auto float_type = mlir::dyn_cast(type); + if (float_type == nullptr) { + return {}; + } + unsigned width = float_type.getWidth(); + if (width != 8 && width != 16) { + return {}; + } + return IntegerType::get(type.getContext(), width); +} + +// Same shape as `type`, with the element type replaced by its storage type, or +// nullptr if the element type is not a narrow float. +RankedTensorType GetStorageTensorType(Type type) { + auto tensor_type = mlir::dyn_cast(type); + if (tensor_type == nullptr) { + return {}; + } + IntegerType storage = GetStorageType(tensor_type.getElementType()); + if (storage == nullptr) { + return {}; + } + return tensor_type.clone(storage); +} + +MemRefType GetStorageMemRefType(MemRefType type) { + IntegerType storage = GetStorageType(type.getElementType()); + if (!storage) { + return nullptr; + } + return MemRefType::Builder(type).setElementType(storage); +} + +// Bitcasts a tensor or scalar `value` to `type`. +Value Bitcast(OpBuilder& builder, Location loc, Type type, Value value) { + if (value.getType() == type) { + return value; + } + if (mlir::isa(type)) { + return tensor::BitcastOp::create(builder, loc, type, value); + } + return ma::BitcastOp::create(builder, loc, type, value); +} + +// Maps `value` from its narrow float type to the storage type. +Value ToStorage(OpBuilder& builder, Location loc, Value value) { + Type type = value.getType(); + if (RankedTensorType tensor_type = GetStorageTensorType(type)) { + return Bitcast(builder, loc, tensor_type, value); + } + if (IntegerType scalar_type = GetStorageType(type)) { + return Bitcast(builder, loc, scalar_type, value); + } + return value; +} + +// op(x : fN) -> bitcast(op(bitcast(x) : iN) : iN -> fN) +template +LogicalResult RetypeDataMovement(OpTy op, PatternRewriter& rewriter) { + Type type = op.getType(); + RankedTensorType storage_type = GetStorageTensorType(type); + if (!storage_type) { + return rewriter.notifyMatchFailure(op, "not a narrow float tensor."); + } + llvm::SmallVector operands = llvm::map_to_vector( + op->getOperands(), + [&](Value operand) { return ToStorage(rewriter, op.getLoc(), operand); }); + auto storage_op = OpTy::create(rewriter, op.getLoc(), storage_type, operands, + op.getProperties(), + op->getDiscardableAttrDictionary().getValue()); + rewriter.replaceOpWithNewOp(op, type, storage_op); + return mlir::success(); +} + +// bitcast(bitcast(x)) -> x, if the types round-trip. +LogicalResult FoldBitcastRoundTrip(tensor::BitcastOp bitcast, + PatternRewriter& rewriter) { + auto producer = bitcast.getSource().getDefiningOp(); + if (!producer || producer.getSource().getType() != bitcast.getType()) { + return rewriter.notifyMatchFailure(bitcast, "not a round trip."); + } + rewriter.replaceOp(bitcast, producer.getSource()); + return mlir::success(); +} + +// Returns a view of narrow float `buffer` as its storage type. The view is +// created right after the definition of `buffer`, so it dominates every use of +// `buffer` and can be shared by all its tile accesses. +Value GetStorageBuffer(PatternRewriter& rewriter, + mlir::TypedValue buffer) { + MemRefType storage_type = GetStorageMemRefType(buffer.getType()); + for (mlir::Operation* user : buffer.getUsers()) { + auto view = mlir::dyn_cast(user); + if (view && view.getType() == storage_type) { + return view; + } + } + OpBuilder::InsertionGuard guard(rewriter); + if (auto arg = mlir::dyn_cast(buffer)) { + rewriter.setInsertionPointToStart(arg.getOwner()); + } else { + rewriter.setInsertionPointAfterValue(buffer); + } + return xtile::MemRefBitcastOp::create(rewriter, buffer.getLoc(), storage_type, + buffer); +} + +// extract(buffer : fN) -> bitcast(extract(view(buffer) : iN) : iN -> fN) +LogicalResult RetypeExtractTile(xtile::ExtractTileOp op, + PatternRewriter& rewriter) { + if (!GetStorageMemRefType(op.getSource().getType())) { + return rewriter.notifyMatchFailure(op, "not a narrow float buffer."); + } + Value buffer = GetStorageBuffer(rewriter, op.getSource()); + auto storage_tile = xtile::ExtractTileOp::create( + rewriter, op.getLoc(), GetStorageTensorType(op.getType()), buffer, + op.getOffsets(), op.getFullTileShapeAttr(), op.getStridesAttr()); + storage_tile->setDiscardableAttrs(op->getDiscardableAttrDictionary()); + rewriter.replaceOpWithNewOp(op, op.getType(), + storage_tile); + return mlir::success(); +} + +// insert(tile : fN, buffer : fN) -> insert(bitcast(tile) : iN, view(buffer)) +LogicalResult RetypeInsertTile(xtile::InsertTileOp op, + PatternRewriter& rewriter) { + if (!GetStorageMemRefType(op.getDestination().getType())) { + return rewriter.notifyMatchFailure(op, "not a narrow float buffer."); + } + Value buffer = GetStorageBuffer(rewriter, op.getDestination()); + Value storage_tile = ToStorage(rewriter, op.getLoc(), op.getSource()); + rewriter.modifyOpInPlace(op, [&] { + op.getSourceMutable().assign(storage_tile); + op.getDestinationMutable().assign(buffer); + }); + return mlir::success(); +} + +class LegalizeNarrowFloatStoragePass + : public impl::LegalizeNarrowFloatStoragePassBase< + LegalizeNarrowFloatStoragePass> { + public: + using LegalizeNarrowFloatStoragePassBase::LegalizeNarrowFloatStoragePassBase; + + void runOnOperation() override { + mlir::MLIRContext* context = &getContext(); + mlir::RewritePatternSet patterns(context); + patterns.add(RetypeDataMovement); + patterns.add(RetypeDataMovement); + patterns.add(RetypeDataMovement); + patterns.add(RetypeDataMovement); + patterns.add(RetypeDataMovement); + patterns.add(RetypeDataMovement); + patterns.add(RetypeDataMovement); + patterns.add(FoldBitcastRoundTrip); + patterns.add(RetypeExtractTile); + patterns.add(RetypeInsertTile); + if (mlir::failed( + mlir::applyPatternsGreedily(getOperation(), std::move(patterns)))) { + signalPassFailure(); + return; + } + } +}; + +} // namespace +} // namespace xla::cpu diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/lower_memref_bitcast.cc b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/lower_memref_bitcast.cc new file mode 100644 index 00000000000000..76863e2a7506e8 --- /dev/null +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/lower_memref_bitcast.cc @@ -0,0 +1,78 @@ +/* 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 "mlir/Conversion/LLVMCommon/LoweringOptions.h" +#include "mlir/Conversion/LLVMCommon/Pattern.h" +#include "mlir/Conversion/LLVMCommon/TypeConverter.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" // IWYU pragma: keep +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/Operation.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Interfaces/DataLayoutInterfaces.h" +#include "mlir/Pass/Pass.h" // IWYU pragma: keep +#include "mlir/Support/LLVM.h" +#include "mlir/Transforms/DialectConversion.h" +#include "xla/backends/cpu/codegen/tiled/transforms/passes.h" // IWYU pragma: keep +#include "xla/codegen/xtile/ir/xtile_ops.h" + +namespace xla::cpu { + +#define GEN_PASS_DEF_LOWERMEMREFBITCASTPASS +#include "xla/backends/cpu/codegen/tiled/transforms/passes.h.inc" + +namespace { + +struct LowerMemRefBitcast + : public mlir::ConvertOpToLLVMPattern { + using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern; + + mlir::LogicalResult matchAndRewrite( + xtile::MemRefBitcastOp op, OpAdaptor adaptor, + mlir::ConversionPatternRewriter& rewriter) const override { + rewriter.replaceOp(op, adaptor.getSource()); + return mlir::success(); + } +}; + +class LowerMemRefBitcastPass + : public impl::LowerMemRefBitcastPassBase { + public: + using LowerMemRefBitcastPassBase::LowerMemRefBitcastPassBase; + + void runOnOperation() override { + mlir::ModuleOp module = getOperation(); + mlir::MLIRContext* context = &getContext(); + mlir::LowerToLLVMOptions options(context, mlir::DataLayout(module)); + mlir::LLVMTypeConverter type_converter(context, options); + + mlir::ConversionTarget target(*context); + target.addIllegalOp(); + target.markUnknownOpDynamicallyLegal([](mlir::Operation*) { return true; }); + + mlir::RewritePatternSet patterns(context); + patterns.add(type_converter); + if (mlir::failed(mlir::applyPartialConversion(module, target, + std::move(patterns)))) { + signalPassFailure(); + return; + } + } +}; + +} // namespace +} // namespace xla::cpu diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/passes.td b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/passes.td index 151d1e3f851c9b..1545a80a5a6007 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/passes.td +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/passes.td @@ -115,6 +115,52 @@ def UnrollVectorsPass : Pass<"xtile-cpu-unroll-vectors"> { } +def LegalizeNarrowFloatStoragePass + : Pass<"xtile-cpu-legalize-narrow-float-storage", "mlir::ModuleOp"> { + let summary = "Moves data movement of narrow floats onto same-width integers."; + + let description = [{ + x86 has no legal vector register class for float types narrower than 32 + bits unless the matching ISA extension is present (AVX512-BF16, + AVX512-FP16, ...). SelectionDAG therefore scalarizes any such vector that + is not immediately extended or truncated, promoting every element to f32 + and paying a libcall (e.g. `__truncsfbf2`) per element on the way back. + This happens for phis, shuffles, masked loads and stores and selects of + narrow float vectors, i.e. for exactly the ops that only move data around. + + This pass treats 8- and 16-bit floats as storage-only types: + + * Tile accesses to narrow float buffers go through an + `xtile.memref_bitcast` view of the buffer as the same-width integer type, + followed by a `tensor.bitcast` back to the float type. + * These bitcasts are then moved through data movement ops (transpose, + reshape, broadcast, slice, concatenate, mask, select) towards the ops that + compute on the float values, e.g. `arith.extf`/`arith.truncf`, where the + remaining pairs cancel. + + Ops the pass does not know simply stop the bitcasts, so it never changes + semantics; at worst the float type survives on that op, as before. + }]; + + let dependentDialects = [ + "::mlir::arith::ArithDialect", + "::mlir::tensor::TensorDialect", + "::xla::xtile::XTileDialect", + ]; +} + +def LowerMemRefBitcastPass + : Pass<"xtile-cpu-lower-memref-bitcast", "mlir::ModuleOp"> { + let summary = "Lowers xtile.memref_bitcast to LLVM."; + + let description = [{ + The memref descriptor holds an opaque pointer, so its LLVM type does not + depend on the element type and the op lowers to nothing. + }]; + + let dependentDialects = ["::mlir::LLVM::LLVMDialect"]; +} + def UnpackSubByteVectorWritePass : Pass<"xtile-cpu-unpack-sub-byte-vector-write"> { let summary = "Unpacks sub-byte vector writes."; diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/legalize_narrow_float_storage.mlir b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/legalize_narrow_float_storage.mlir new file mode 100644 index 00000000000000..093fcfacba9e86 --- /dev/null +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/legalize_narrow_float_storage.mlir @@ -0,0 +1,153 @@ +// Copyright 2026 The OpenXLA Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// ============================================================================== + +// RUN: fusion_compiler_opt %s -xtile-cpu-legalize-narrow-float-storage \ +// RUN: -split-input-file | FileCheck %s + +func.func @extract_extf(%arg0: memref<16xbf16>, %arg1: index) -> tensor<8xf32> { + %0 = xtile.extract %arg0[%arg1] [8] [1] : memref<16xbf16> -> tensor<8xbf16> + %1 = arith.extf %0 : tensor<8xbf16> to tensor<8xf32> + return %1 : tensor<8xf32> +} +// CHECK-LABEL: func.func @extract_extf( +// CHECK-SAME: %[[ARG0:.*]]: memref<16xbf16>, %[[ARG1:.*]]: index) +// CHECK: %[[VIEW:.*]] = xtile.memref_bitcast %[[ARG0]] : memref<16xbf16> -> memref<16xi16> +// CHECK: %[[TILE:.*]] = xtile.extract %[[VIEW]][%[[ARG1]]] [8] [1] : memref<16xi16> -> tensor<8xi16> +// CHECK: %[[CAST:.*]] = tensor.bitcast %[[TILE]] : tensor<8xi16> to tensor<8xbf16> +// CHECK: arith.extf %[[CAST]] : tensor<8xbf16> to tensor<8xf32> + +// ----- + +func.func @copy(%arg0: memref<16xbf16>, %arg1: memref<16xbf16>, %arg2: index) { + %0 = xtile.extract %arg0[%arg2] [8] [1] : memref<16xbf16> -> tensor<8xbf16> + xtile.insert %0 into %arg1[%arg2] [8] [1] : tensor<8xbf16> -> memref<16xbf16> + return +} +// CHECK-LABEL: func.func @copy( +// CHECK-SAME: %[[ARG0:.*]]: memref<16xbf16>, %[[ARG1:.*]]: memref<16xbf16>, %[[ARG2:.*]]: index) +// CHECK-DAG: %[[SRC:.*]] = xtile.memref_bitcast %[[ARG0]] : memref<16xbf16> -> memref<16xi16> +// CHECK-DAG: %[[DST:.*]] = xtile.memref_bitcast %[[ARG1]] : memref<16xbf16> -> memref<16xi16> +// CHECK: %[[TILE:.*]] = xtile.extract %[[SRC]][%[[ARG2]]] [8] [1] : memref<16xi16> -> tensor<8xi16> +// CHECK-NOT: tensor.bitcast +// CHECK: xtile.insert %[[TILE]] into %[[DST]][%[[ARG2]]] [8] [1] : tensor<8xi16> -> memref<16xi16> + +// ----- + +// All tile accesses to a buffer share one view of it. +func.func @shared_view(%arg0: memref<16xf16>, %arg1: index, %arg2: index) -> tensor<8xf32> { + %0 = xtile.extract %arg0[%arg1] [8] [1] : memref<16xf16> -> tensor<8xf16> + %1 = xtile.extract %arg0[%arg2] [8] [1] : memref<16xf16> -> tensor<8xf16> + %2 = arith.extf %0 : tensor<8xf16> to tensor<8xf32> + %3 = arith.extf %1 : tensor<8xf16> to tensor<8xf32> + %4 = arith.addf %2, %3 : tensor<8xf32> + return %4 : tensor<8xf32> +} +// CHECK-LABEL: func.func @shared_view( +// CHECK: %[[VIEW:.*]] = xtile.memref_bitcast %{{.*}} : memref<16xf16> -> memref<16xi16> +// CHECK-NOT: xtile.memref_bitcast +// CHECK: xtile.extract %[[VIEW]] +// CHECK: xtile.extract %[[VIEW]] + +// ----- + +// The transpose of the truncated tile runs on the storage type. +func.func @truncf_transpose(%arg0: memref<8x8xf32>, %arg1: memref<8x8xbf16>, %arg2: index) { + %0 = xtile.extract %arg0[%arg2, %arg2] [8, 8] [1, 1] : memref<8x8xf32> -> tensor<8x8xf32> + %1 = arith.truncf %0 : tensor<8x8xf32> to tensor<8x8xbf16> + %2 = stablehlo.transpose %1, dims = [1, 0] : (tensor<8x8xbf16>) -> tensor<8x8xbf16> + xtile.insert %2 into %arg1[%arg2, %arg2] [8, 8] [1, 1] : tensor<8x8xbf16> -> memref<8x8xbf16> + return +} +// CHECK-LABEL: func.func @truncf_transpose( +// CHECK-SAME: %[[ARG0:.*]]: memref<8x8xf32>, %[[ARG1:.*]]: memref<8x8xbf16>, %[[ARG2:.*]]: index) +// CHECK: %[[DST:.*]] = xtile.memref_bitcast %[[ARG1]] : memref<8x8xbf16> -> memref<8x8xi16> +// CHECK: %[[TILE:.*]] = xtile.extract %[[ARG0]] +// CHECK: %[[TRUNC:.*]] = arith.truncf %[[TILE]] : tensor<8x8xf32> to tensor<8x8xbf16> +// CHECK: %[[CAST:.*]] = tensor.bitcast %[[TRUNC]] : tensor<8x8xbf16> to tensor<8x8xi16> +// CHECK: %[[TRANS:.*]] = stablehlo.transpose %[[CAST]], dims = [1, 0] : (tensor<8x8xi16>) -> tensor<8x8xi16> +// CHECK: xtile.insert %[[TRANS]] into %[[DST]] + +// ----- + +// A chain of data movement ops between tile accesses runs entirely on the +// storage type, without any bitcasts in between. +func.func @data_movement_chain(%arg0: memref<4x8xf8E4M3FN>, %arg1: memref<8x8xf8E4M3FN>, %arg2: index) { + %0 = xtile.extract %arg0[%arg2, %arg2] [4, 8] [1, 1] : memref<4x8xf8E4M3FN> -> tensor<4x8xf8E4M3FN> + %1 = stablehlo.concatenate %0, %0, dim = 0 : (tensor<4x8xf8E4M3FN>, tensor<4x8xf8E4M3FN>) -> tensor<8x8xf8E4M3FN> + %2 = stablehlo.transpose %1, dims = [1, 0] : (tensor<8x8xf8E4M3FN>) -> tensor<8x8xf8E4M3FN> + %3 = stablehlo.reshape %2 : (tensor<8x8xf8E4M3FN>) -> tensor<64xf8E4M3FN> + %4 = stablehlo.slice %3 [0:32] : (tensor<64xf8E4M3FN>) -> tensor<32xf8E4M3FN> + %5 = stablehlo.broadcast_in_dim %4, dims = [1] : (tensor<32xf8E4M3FN>) -> tensor<2x32xf8E4M3FN> + %6 = stablehlo.reshape %5 : (tensor<2x32xf8E4M3FN>) -> tensor<8x8xf8E4M3FN> + xtile.insert %6 into %arg1[%arg2, %arg2] [8, 8] [1, 1] : tensor<8x8xf8E4M3FN> -> memref<8x8xf8E4M3FN> + return +} +// CHECK-LABEL: func.func @data_movement_chain( +// CHECK-NOT: tensor.bitcast +// CHECK: %[[TILE:.*]] = xtile.extract {{.*}} : memref<4x8xi8> -> tensor<4x8xi8> +// CHECK: %[[CONCAT:.*]] = stablehlo.concatenate %[[TILE]], %[[TILE]], dim = 0 : (tensor<4x8xi8>, tensor<4x8xi8>) -> tensor<8x8xi8> +// CHECK: %[[TRANS:.*]] = stablehlo.transpose %[[CONCAT]], dims = [1, 0] : (tensor<8x8xi8>) -> tensor<8x8xi8> +// CHECK: %[[RESHAPE:.*]] = stablehlo.reshape %[[TRANS]] : (tensor<8x8xi8>) -> tensor<64xi8> +// CHECK: %[[SLICE:.*]] = stablehlo.slice %[[RESHAPE]] [0:32] : (tensor<64xi8>) -> tensor<32xi8> +// CHECK: %[[BCAST:.*]] = stablehlo.broadcast_in_dim %[[SLICE]], dims = [1] : (tensor<32xi8>) -> tensor<2x32xi8> +// CHECK: %[[RESHAPE2:.*]] = stablehlo.reshape %[[BCAST]] : (tensor<2x32xi8>) -> tensor<8x8xi8> +// CHECK: xtile.insert %[[RESHAPE2]] into {{.*}} : tensor<8x8xi8> -> memref<8x8xi8> +// CHECK-NOT: tensor.bitcast + +// ----- + +func.func @select(%arg0: memref<16xbf16>, %arg1: memref<16xbf16>, %arg2: tensor<8xi1>, %arg3: index) { + %0 = xtile.extract %arg0[%arg3] [8] [1] : memref<16xbf16> -> tensor<8xbf16> + %1 = xtile.extract %arg1[%arg3] [8] [1] : memref<16xbf16> -> tensor<8xbf16> + %2 = arith.select %arg2, %0, %1 : tensor<8xi1>, tensor<8xbf16> + xtile.insert %2 into %arg0[%arg3] [8] [1] : tensor<8xbf16> -> memref<16xbf16> + return +} +// CHECK-LABEL: func.func @select( +// CHECK-SAME: %[[ARG0:.*]]: memref<16xbf16>, %[[ARG1:.*]]: memref<16xbf16>, %[[COND:.*]]: tensor<8xi1>, %[[ARG3:.*]]: index) +// CHECK-DAG: %[[VIEW0:.*]] = xtile.memref_bitcast %[[ARG0]] +// CHECK-DAG: %[[VIEW1:.*]] = xtile.memref_bitcast %[[ARG1]] +// CHECK: %[[LHS:.*]] = xtile.extract %[[VIEW0]] +// CHECK: %[[RHS:.*]] = xtile.extract %[[VIEW1]] +// CHECK: %[[SEL:.*]] = arith.select %[[COND]], %[[LHS]], %[[RHS]] : tensor<8xi1>, tensor<8xi16> +// CHECK: xtile.insert %[[SEL]] into %[[VIEW0]] + +// ----- + +func.func @mask(%arg0: memref<16xbf16>, %arg1: memref<16xbf16>, %arg2: index) { + %cst = arith.constant 0.0 : bf16 + %0 = xtile.extract %arg0[%arg2] [8] [1] : memref<16xbf16> -> tensor<8xbf16> + %1 = xtile.mask %0 bounds [5], %cst : tensor<8xbf16> + xtile.insert %1 into %arg1[%arg2] [8] [1] : tensor<8xbf16> -> memref<16xbf16> + return +} +// CHECK-LABEL: func.func @mask( +// CHECK: %[[ZERO:.*]] = arith.constant 0 : i16 +// CHECK: %[[TILE:.*]] = xtile.extract {{.*}} : memref<16xi16> -> tensor<8xi16> +// CHECK: %[[MASKED:.*]] = xtile.mask %[[TILE]] bounds [5], %[[ZERO]] : tensor<8xi16> +// CHECK: xtile.insert %[[MASKED]] into {{.*}} : tensor<8xi16> -> memref<16xi16> + +// ----- + +// Types that are not 8- or 16-bit floats are left alone. +func.func @not_narrow(%arg0: memref<16xf32>, %arg1: memref<16xf4E2M1FN>, %arg2: index) -> (tensor<8xf32>, tensor<8xf4E2M1FN>) { + %0 = xtile.extract %arg0[%arg2] [8] [1] : memref<16xf32> -> tensor<8xf32> + %1 = xtile.extract %arg1[%arg2] [8] [1] : memref<16xf4E2M1FN> -> tensor<8xf4E2M1FN> + return %0, %1 : tensor<8xf32>, tensor<8xf4E2M1FN> +} +// CHECK-LABEL: func.func @not_narrow( +// CHECK-NOT: xtile.memref_bitcast +// CHECK: xtile.extract %{{.*}} : memref<16xf32> -> tensor<8xf32> +// CHECK: xtile.extract %{{.*}} : memref<16xf4E2M1FN> -> tensor<8xf4E2M1FN> diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/lower_memref_bitcast.mlir b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/lower_memref_bitcast.mlir new file mode 100644 index 00000000000000..7dbd101c3cf36c --- /dev/null +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/lower_memref_bitcast.mlir @@ -0,0 +1,43 @@ +// Copyright 2026 The OpenXLA Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// ============================================================================== + +// RUN: fusion_compiler_opt %s -xtile-cpu-lower-memref-bitcast \ +// RUN: -split-input-file | FileCheck %s + +func.func @memref_bitcast(%arg0: memref<16xbf16>, %arg1: index) -> vector<8xi16> { + %c0 = arith.constant 0 : i16 + %0 = xtile.memref_bitcast %arg0 : memref<16xbf16> -> memref<16xi16> + %1 = vector.transfer_read %0[%arg1], %c0 : memref<16xi16>, vector<8xi16> + return %1 : vector<8xi16> +} +// CHECK-LABEL: func.func @memref_bitcast( +// CHECK-SAME: %[[ARG0:.*]]: memref<16xbf16>, %[[ARG1:.*]]: index) +// CHECK-NOT: xtile.memref_bitcast +// CHECK: %[[DESC:.*]] = builtin.unrealized_conversion_cast %[[ARG0]] : memref<16xbf16> to !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> +// CHECK: %[[VIEW:.*]] = builtin.unrealized_conversion_cast %[[DESC]] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> to memref<16xi16> +// CHECK: vector.transfer_read %[[VIEW]][%[[ARG1]]] + +// ----- + +func.func @memref_bitcast_2d(%arg0: memref<4x8xf8E4M3FN>) -> memref<4x8xi8> { + %0 = xtile.memref_bitcast %arg0 : memref<4x8xf8E4M3FN> -> memref<4x8xi8> + return %0 : memref<4x8xi8> +} +// CHECK-LABEL: func.func @memref_bitcast_2d( +// CHECK-SAME: %[[ARG0:.*]]: memref<4x8xf8E4M3FN>) +// CHECK-NOT: xtile.memref_bitcast +// CHECK: %[[DESC:.*]] = builtin.unrealized_conversion_cast %[[ARG0]] : memref<4x8xf8E4M3FN> to !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)> +// CHECK: %[[VIEW:.*]] = builtin.unrealized_conversion_cast %[[DESC]] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)> to memref<4x8xi8> +// CHECK: return %[[VIEW]] diff --git a/third_party/xla/xla/backends/gpu/codegen/BUILD b/third_party/xla/xla/backends/gpu/codegen/BUILD index 0877e5e0921bf2..52bfcd607cb8e3 100644 --- a/third_party/xla/xla/backends/gpu/codegen/BUILD +++ b/third_party/xla/xla/backends/gpu/codegen/BUILD @@ -112,6 +112,50 @@ xla_test( ], ) +xla_test( + name = "cudnn_scaled_dot_test", + srcs = ["cudnn_scaled_dot_test.cc"], + backends = ["b200"], + tags = ["cuda-only"], + deps = [ + "//xla:error_spec", + "//xla:literal", + "//xla:shape_util", + "//xla:types", + "//xla:xla_data_proto_cc", + "//xla:xla_proto_cc", + "//xla/backends/autotuner:backends_proto_cc", + "//xla/backends/gpu/tests:hlo_pjrt_gpu_test_base", + "//xla/backends/gpu/transforms:composite_rewriter", + "//xla/backends/gpu/transforms:cudnn_custom_call_compiler", + "//xla/backends/gpu/transforms:cudnn_fusion_compiler", + "//xla/backends/gpu/transforms:scaled_dot_rewriter", + "//xla/hlo/ir:hlo", + "//xla/hlo/parser:hlo_parser", + "//xla/hlo/testlib:verified_hlo_module", + "//xla/service/gpu:backend_configs_cc", + "//xla/service/gpu:ir_emission_utils", + "//xla/stream_executor:device_description", + "//xla/stream_executor:platform_manager", + "//xla/stream_executor:semantic_version", + "//xla/stream_executor:stream_executor_h", + "//xla/stream_executor/cuda:cuda_compute_capability", + "//xla/tests:hlo_interpreter_reference_mixin", + "//xla/tests:test_utils", + "//xla/tsl/platform:test", + "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/log", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/random:distributions", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_macros", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/types:span", + "@com_google_googletest//:gtest_main", + ], +) + cc_library( name = "custom", srcs = ["custom.cc"], diff --git a/third_party/xla/xla/backends/gpu/codegen/cudnn_scaled_dot_test.cc b/third_party/xla/xla/backends/gpu/codegen/cudnn_scaled_dot_test.cc new file mode 100644 index 00000000000000..8a9a98e060af43 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/codegen/cudnn_scaled_dot_test.cc @@ -0,0 +1,570 @@ +/* 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 +#include +#include +#include +#include +#include +#include + +#include +#include +#include "absl/algorithm/container.h" +#include "absl/log/check.h" +#include "absl/log/log.h" +#include "absl/random/uniform_int_distribution.h" +#include "absl/status/status.h" +#include "absl/status/status_macros.h" +#include "absl/status/statusor.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/str_replace.h" +#include "absl/strings/string_view.h" +#include "absl/types/span.h" +#include "xla/backends/autotuner/backends.pb.h" +#include "xla/backends/gpu/tests/hlo_pjrt_gpu_test_base.h" +#include "xla/backends/gpu/transforms/composite_rewriter.h" +#include "xla/backends/gpu/transforms/cudnn_custom_call_compiler.h" +#include "xla/backends/gpu/transforms/cudnn_fusion_compiler.h" +#include "xla/backends/gpu/transforms/scaled_dot_rewriter.h" +#include "xla/error_spec.h" +#include "xla/hlo/ir/hlo_computation.h" +#include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/parser/hlo_parser.h" +#include "xla/hlo/testlib/verified_hlo_module.h" +#include "xla/literal.h" +#include "xla/primitive_util.h" +#include "xla/service/gpu/backend_configs.pb.h" +#include "xla/service/gpu/ir_emission_utils.h" +#include "xla/stream_executor/cuda/cuda_compute_capability.h" +#include "xla/stream_executor/device_description.h" +#include "xla/stream_executor/platform_manager.h" +#include "xla/stream_executor/semantic_version.h" +#include "xla/stream_executor/stream_executor.h" +#include "xla/tests/hlo_interpreter_reference_mixin.h" +#include "xla/tests/test_utils.h" +#include "xla/tsl/platform/test.h" +#include "xla/types.h" +#include "xla/xla.pb.h" +#include "xla/xla_data.pb.h" + +namespace xla { +namespace gpu { +namespace { + +class CudnnScaledDotTestBase + : public HloInterpreterReferenceMixin { + public: + se::StreamExecutor* stream_executor() const { + auto platform = + se::PlatformManager::PlatformWithId(stream_executor_platform_id()); + CHECK_OK(platform); + auto executor = (*platform)->ExecutorForDevice(0); + CHECK_OK(executor); + return *executor; + } + + DebugOptions GetDebugOptionsForTest() const override { + DebugOptions debug_options = HloPjRtGpuTestBase::GetDebugOptionsForTest(); + debug_options.set_xla_gpu_autotune_level(0); + debug_options.set_xla_gpu_cudnn_gemm_fusion_level(2); + debug_options.clear_xla_gpu_experimental_autotune_backends(); + debug_options.add_xla_gpu_experimental_autotune_backends( + autotuner::Backend::CUDNN); + debug_options.add_xla_gpu_experimental_autotune_backends( + autotuner::Backend::CUBLASLT); + debug_options.set_xla_gpu_experimental_scaled_dot_with_triton(false); + return debug_options; + } + + se::CudaComputeCapability get_cuda_cc() const { + return device_description().cuda_compute_capability(); + } + + bool IsAtLeastBlackwellWithCuDnn9() const { + return get_cuda_cc().IsAtLeastBlackwell() && + gpu_target_config() + .device_description.dnn_version() + .major_version() >= 9; + } + + absl::Status PopulateScale(Literal* scale, std::minstd_rand0* engine) { + switch (scale->shape().element_type()) { + case F8E8M0FNU: { + absl::uniform_int_distribution exponent_distribution(-4, 1); + return scale->Populate( + [&](absl::Span /*indices*/) { + return float8_e8m0fnu( + std::ldexp(1.0f, exponent_distribution(*engine))); + }); + } + case F8E4M3FN: + for (float8_e4m3fn& value : scale->data()) { + value = float8_e4m3fn(std::abs(static_cast(value))); + } + return absl::OkStatus(); + case F8E5M2: + for (float8_e5m2& value : scale->data()) { + value = float8_e5m2(std::abs(static_cast(value))); + } + return absl::OkStatus(); + default: + return absl::InvalidArgumentError( + absl::StrCat("Unsupported scale type: ", + PrimitiveType_Name(scale->shape().element_type()))); + } + } + + absl::StatusOr> MakeScaledDotArguments( + const HloModule* module) { + std::minstd_rand0 engine; + FakeArgumentsOptions options; + options.engine = &engine; + ABSL_ASSIGN_OR_RETURN(std::vector arguments, + MakeFakeArguments(module, options)); + if (arguments.size() < 4) { + return absl::InternalError(absl::StrCat( + "Expected at least 4 scaled-dot arguments, got ", arguments.size())); + } + ABSL_RETURN_IF_ERROR(PopulateScale(&arguments[2], &engine)); + ABSL_RETURN_IF_ERROR(PopulateScale(&arguments[3], &engine)); + return arguments; + } +}; + +enum class CuDnnSupportResult { + kSupported, // Lowered to native cuDNN block-scaled graph, execution succeeds + kNotSupported, // Not supported by cuDNN block scaling, falls back to + // decomposed dot + kFailed, // Lowered to native cuDNN block-scaled graph, but execution + // exhibits discrepancy +}; + +absl::string_view CuDnnSupportResultToString(CuDnnSupportResult result) { + switch (result) { + case CuDnnSupportResult::kSupported: + return "supported"; + case CuDnnSupportResult::kNotSupported: + return "not_supported"; + case CuDnnSupportResult::kFailed: + return "failed"; + } +} + +std::ostream& operator<<(std::ostream& os, CuDnnSupportResult result) { + return os << CuDnnSupportResultToString(result); +} + +struct ScaledDotCoverageTestCase { + PrimitiveType lhs_type; + PrimitiveType rhs_type; + PrimitiveType scale_type; + int block_size; + bool lhs_k_minor; + bool rhs_k_minor; + CuDnnSupportResult expected_result; +}; + +std::vector GetCoverageTestCases() { + return { + // 1. MXFP8 E4M3FN x E4M3FN (scale E8M0FNU, block 32) + {F8E4M3FN, F8E4M3FN, F8E8M0FNU, 32, false, false, + CuDnnSupportResult::kNotSupported}, + {F8E4M3FN, F8E4M3FN, F8E8M0FNU, 32, false, true, + CuDnnSupportResult::kNotSupported}, + {F8E4M3FN, F8E4M3FN, F8E8M0FNU, 32, true, false, + CuDnnSupportResult::kNotSupported}, + {F8E4M3FN, F8E4M3FN, F8E8M0FNU, 32, true, true, + CuDnnSupportResult::kSupported}, + + // 2. MXFP8 E4M3FN x E5M2 (scale E8M0FNU, block 32) + {F8E4M3FN, F8E5M2, F8E8M0FNU, 32, false, false, + CuDnnSupportResult::kNotSupported}, + {F8E4M3FN, F8E5M2, F8E8M0FNU, 32, false, true, + CuDnnSupportResult::kNotSupported}, + {F8E4M3FN, F8E5M2, F8E8M0FNU, 32, true, false, + CuDnnSupportResult::kNotSupported}, + {F8E4M3FN, F8E5M2, F8E8M0FNU, 32, true, true, + CuDnnSupportResult::kSupported}, + + // 3. E4M3FN x F4E2M1FN (scale E8M0FNU, block 32) + {F8E4M3FN, F4E2M1FN, F8E8M0FNU, 32, false, false, + CuDnnSupportResult::kNotSupported}, + {F8E4M3FN, F4E2M1FN, F8E8M0FNU, 32, false, true, + CuDnnSupportResult::kNotSupported}, + {F8E4M3FN, F4E2M1FN, F8E8M0FNU, 32, true, false, + CuDnnSupportResult::kNotSupported}, + {F8E4M3FN, F4E2M1FN, F8E8M0FNU, 32, true, true, + CuDnnSupportResult::kNotSupported}, + + // 4. MXFP8 E5M2 x E4M3FN (scale E8M0FNU, block 32) + {F8E5M2, F8E4M3FN, F8E8M0FNU, 32, false, false, + CuDnnSupportResult::kNotSupported}, + {F8E5M2, F8E4M3FN, F8E8M0FNU, 32, false, true, + CuDnnSupportResult::kNotSupported}, + {F8E5M2, F8E4M3FN, F8E8M0FNU, 32, true, false, + CuDnnSupportResult::kNotSupported}, + {F8E5M2, F8E4M3FN, F8E8M0FNU, 32, true, true, + CuDnnSupportResult::kSupported}, + + // 5. MXFP8 E5M2 x E5M2 (scale E8M0FNU, block 32) - unsupported by cuDNN + {F8E5M2, F8E5M2, F8E8M0FNU, 32, false, false, + CuDnnSupportResult::kNotSupported}, + {F8E5M2, F8E5M2, F8E8M0FNU, 32, false, true, + CuDnnSupportResult::kNotSupported}, + {F8E5M2, F8E5M2, F8E8M0FNU, 32, true, false, + CuDnnSupportResult::kNotSupported}, + {F8E5M2, F8E5M2, F8E8M0FNU, 32, true, true, + CuDnnSupportResult::kNotSupported}, + + // 6. E5M2 x F4E2M1FN (scale E8M0FNU, block 32) + {F8E5M2, F4E2M1FN, F8E8M0FNU, 32, false, false, + CuDnnSupportResult::kNotSupported}, + {F8E5M2, F4E2M1FN, F8E8M0FNU, 32, false, true, + CuDnnSupportResult::kNotSupported}, + {F8E5M2, F4E2M1FN, F8E8M0FNU, 32, true, false, + CuDnnSupportResult::kNotSupported}, + {F8E5M2, F4E2M1FN, F8E8M0FNU, 32, true, true, + CuDnnSupportResult::kNotSupported}, + + // 7. F4E2M1FN x E4M3FN (scale E8M0FNU, block 32) + {F4E2M1FN, F8E4M3FN, F8E8M0FNU, 32, false, false, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F8E4M3FN, F8E8M0FNU, 32, false, true, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F8E4M3FN, F8E8M0FNU, 32, true, false, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F8E4M3FN, F8E8M0FNU, 32, true, true, + CuDnnSupportResult::kNotSupported}, + + // 8. F4E2M1FN x E5M2 (scale E8M0FNU, block 32) + {F4E2M1FN, F8E5M2, F8E8M0FNU, 32, false, false, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F8E5M2, F8E8M0FNU, 32, false, true, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F8E5M2, F8E8M0FNU, 32, true, false, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F8E5M2, F8E8M0FNU, 32, true, true, + CuDnnSupportResult::kNotSupported}, + + // 9. F4E2M1FN x F4E2M1FN (scale E8M0FNU, block 32) + {F4E2M1FN, F4E2M1FN, F8E8M0FNU, 32, false, false, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F4E2M1FN, F8E8M0FNU, 32, false, true, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F4E2M1FN, F8E8M0FNU, 32, true, false, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F4E2M1FN, F8E8M0FNU, 32, true, true, + CuDnnSupportResult::kNotSupported}, + + // 10. F4E2M1FN x F4E2M1FN (scale E8M0FNU, block 16) + {F4E2M1FN, F4E2M1FN, F8E8M0FNU, 16, false, false, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F4E2M1FN, F8E8M0FNU, 16, false, true, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F4E2M1FN, F8E8M0FNU, 16, true, false, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F4E2M1FN, F8E8M0FNU, 16, true, true, + CuDnnSupportResult::kNotSupported}, + + // 11. NVFP4 F4E2M1FN x F4E2M1FN (scale E4M3FN, block 16) + {F4E2M1FN, F4E2M1FN, F8E4M3FN, 16, false, false, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F4E2M1FN, F8E4M3FN, 16, false, true, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F4E2M1FN, F8E4M3FN, 16, true, false, + CuDnnSupportResult::kNotSupported}, + {F4E2M1FN, F4E2M1FN, F8E4M3FN, 16, true, true, + CuDnnSupportResult::kFailed}, + }; +} + +std::string ScaledDotCoverageTestParamToString( + const ::testing::TestParamInfo& info) { + const auto& tc = info.param; + return absl::StrCat( + PrimitiveType_Name(tc.lhs_type), "_", PrimitiveType_Name(tc.rhs_type), + "_", PrimitiveType_Name(tc.scale_type), "_Block", tc.block_size, "_Lhs", + tc.lhs_k_minor, "_Rhs", tc.rhs_k_minor); +} + +class CudnnScaledDotCoverageTest + : public CudnnScaledDotTestBase, + public ::testing::WithParamInterface {}; + +TEST_P(CudnnScaledDotCoverageTest, Executes) { + const auto& param = GetParam(); + const bool lhs_k_minor = param.lhs_k_minor; + const bool rhs_k_minor = param.rhs_k_minor; + + auto cc = get_cuda_cc(); + std::string device_name = "Unknown"; + if (cc.IsAtLeastBlackwell()) { + device_name = "B200"; + } else if (cc.IsAtLeastHopper()) { + device_name = "H100"; + } else { + device_name = "Pre-Hopper"; + } + + std::string lhs_name = + primitive_util::LowercasePrimitiveTypeName(param.lhs_type); + std::string rhs_name = + primitive_util::LowercasePrimitiveTypeName(param.rhs_type); + std::string scale_name = + primitive_util::LowercasePrimitiveTypeName(param.scale_type); + std::string lhs_display = + lhs_k_minor ? lhs_name : absl::StrCat(lhs_name, ".T"); + std::string rhs_display = + rhs_k_minor ? absl::StrCat(rhs_name, ".T") : rhs_name; + + LOG(ERROR) << "Report Device: " << device_name; + LOG(ERROR) << "Report LHS: " << lhs_display; + LOG(ERROR) << "Report RHS: " << rhs_display; + LOG(ERROR) << "Report Scale: " << scale_name; + LOG(ERROR) << "Report BlockSize: " << param.block_size; + + constexpr int64_t m = 128; + constexpr int64_t n = 128; + constexpr int64_t k = 256; + const int64_t scale_k = k / param.block_size; + + const std::string lhs_shape = + lhs_k_minor ? absl::StrCat(m, ",", k) : absl::StrCat(k, ",", m); + const std::string rhs_shape = + rhs_k_minor ? absl::StrCat(n, ",", k) : absl::StrCat(k, ",", n); + const std::string lhs_scale_shape = lhs_k_minor + ? absl::StrCat(m, ",", scale_k) + : absl::StrCat(scale_k, ",", m); + const std::string rhs_scale_shape = rhs_k_minor + ? absl::StrCat(n, ",", scale_k) + : absl::StrCat(scale_k, ",", n); + + // 1. Stage: After Composite + std::string after_composite = "N/A"; + constexpr absl::string_view kCompositeHloTemplate = R"hlo( +HloModule test_composite + +%xla.scaled_dot.1 { + %p0 = $lhs_type[$lhs_shape]{1,0} parameter(0) + %p1 = $rhs_type[$rhs_shape]{1,0} parameter(1) + %p2 = $scale_type[$lhs_scale_shape]{1,0} parameter(2) + %p3 = $scale_type[$rhs_scale_shape]{1,0} parameter(3) + %zero = bf16[] constant(0) + ROOT %dummy = bf16[$output_shape]{1,0} broadcast(%zero), dimensions={} +} + +ENTRY %main { + %lhs = $lhs_type[$lhs_shape]{1,0} parameter(0) + %rhs = $rhs_type[$rhs_shape]{1,0} parameter(1) + %lhs_scale = $scale_type[$lhs_scale_shape]{1,0} parameter(2) + %rhs_scale = $scale_type[$rhs_scale_shape]{1,0} parameter(3) + ROOT %call = bf16[$output_shape]{1,0} call(%lhs, %rhs, %lhs_scale, %rhs_scale), + to_apply=%xla.scaled_dot.1, + is_composite=true, + frontend_attributes={ + composite.attributes="{dimension_numbers=[[[$lhs_contracting_dim],[$rhs_contracting_dim]],[[],[]]]}", + composite.name="xla.scaled_dot", + composite.version="1" + } +} +)hlo"; + + std::string composite_hlo = + absl::StrReplaceAll(kCompositeHloTemplate, + {{"$lhs_type", lhs_name}, + {"$rhs_type", rhs_name}, + {"$scale_type", scale_name}, + {"$lhs_shape", lhs_shape}, + {"$rhs_shape", rhs_shape}, + {"$lhs_scale_shape", lhs_scale_shape}, + {"$rhs_scale_shape", rhs_scale_shape}, + {"$output_shape", absl::StrCat(m, ",", n)}, + {"$lhs_contracting_dim", lhs_k_minor ? "1" : "0"}, + {"$rhs_contracting_dim", rhs_k_minor ? "1" : "0"}}); + + if (auto comp_module = ParseAndReturnVerifiedModule(composite_hlo); + comp_module.ok()) { + CompositeRewriter rewriter; + if (auto rewrite_status = rewriter.Run(comp_module->get()); + rewrite_status.ok()) { + const auto* root = + (*comp_module)->entry_computation()->root_instruction(); + after_composite = HloOpcodeString(root->opcode()); + } + } + LOG(ERROR) << "Report AfterComposite: " << after_composite; + EXPECT_EQ(after_composite, "scaled-dot"); + + // 2. Stage: Optimized HLO (compiled directly from composite_hlo) + std::string optimized_hlo = "N/A"; + HloInstruction* custom_call_or_fusion = nullptr; + auto optimized_module_or = GetOptimizedModule(composite_hlo); + if (!optimized_module_or.ok()) { + LOG(ERROR) << "GetOptimizedModule error: " << optimized_module_or.status(); + } + if (optimized_module_or.ok()) { + HloModule& optimized_module = **optimized_module_or; + for (HloComputation* comp : optimized_module.computations()) { + for (HloInstruction* instr : comp->instructions()) { + if (instr->opcode() == HloOpcode::kCustomCall && + instr->custom_call_target() == "__cudnn$blockScaledDot") { + optimized_hlo = "__cudnn$blockScaledDot"; + custom_call_or_fusion = instr; + break; + } + if (instr->opcode() == HloOpcode::kFusion) { + if (auto gpu_cfg = instr->backend_config(); + gpu_cfg.ok() && + gpu_cfg->fusion_backend_config().kind() == "__cudnn$fusion") { + const HloComputation* fused_comp = + instr->fused_instructions_computation(); + bool has_scaled_dot = false; + for (const HloInstruction* fused_instr : + fused_comp->instructions()) { + if (fused_instr->opcode() == HloOpcode::kScaledDot) { + has_scaled_dot = true; + break; + } + } + optimized_hlo = has_scaled_dot ? "__cudnn$fusion" : "dot"; + custom_call_or_fusion = instr; + break; + } + } + if (instr->opcode() == HloOpcode::kScaledDot) { + optimized_hlo = "scaled-dot"; + } else if (instr->opcode() == HloOpcode::kDot && + optimized_hlo == "N/A") { + optimized_hlo = "dot"; + } + } + if (custom_call_or_fusion != nullptr) { + break; + } + } + } + LOG(ERROR) << "Report OptimizedHLO: " << optimized_hlo; + + // 3. Stage: cuDNN Graph + std::string cudnn_graph_pattern = "N/A"; + if (custom_call_or_fusion != nullptr && IsAtLeastBlackwellWithCuDnn9() && + optimized_hlo != "dot" && optimized_module_or.ok()) { + HloModule& optimized_module = **optimized_module_or; + if (custom_call_or_fusion->opcode() == HloOpcode::kFusion) { + BinaryMap compilation_results; + CuDnnFusionCompiler cudnn_compiler(stream_executor()->AsDnn(), + device_description(), + compilation_results); + auto cloned = optimized_module.Clone(); + auto res = cudnn_compiler.Run(cloned.get()); + if (res.ok()) { + cudnn_graph_pattern = "BLOCK_SCALE_DEQUANTIZE + MATMUL"; + } else { + cudnn_graph_pattern = absl::StrCat("FAILED: ", res.status().message()); + } + } else if (custom_call_or_fusion->opcode() == HloOpcode::kCustomCall) { + BinaryMap compilation_results; + CuDnnCustomCallCompiler custom_call_compiler(stream_executor()->AsDnn(), + device_description(), + compilation_results); + auto cloned = optimized_module.Clone(); + auto res = custom_call_compiler.Run(cloned.get()); + if (res.ok() && !compilation_results.empty()) { + cudnn_graph_pattern = "BLOCK_SCALE_DEQUANTIZE + MATMUL"; + } else if (!res.ok()) { + cudnn_graph_pattern = absl::StrCat("FAILED: ", res.status().message()); + } + } + } + + // 4. Stage: Execution Status + std::string status = "FAILURE"; + if (!IsAtLeastBlackwellWithCuDnn9()) { + GTEST_SKIP() + << "Block scaled dot cuDNN execution requires Blackwell and cuDNN 9+."; + } + if (optimized_module_or.ok()) { + auto ref_module_or = ParseAndReturnVerifiedModule(composite_hlo); + if (ref_module_or.ok()) { + CompositeRewriter composite_pass; + (void)composite_pass.Run(ref_module_or->get()); + ScaledDotRewriter decompose_pass; + if (auto decomposed = decompose_pass.Run(ref_module_or->get()); + decomposed.ok()) { + auto ref_optimized_or = + GetOptimizedModule((*ref_module_or)->ToString()); + if (ref_optimized_or.ok()) { + auto run_result = RunAndCompareTwoModules( + std::move(*optimized_module_or), std::move(*ref_optimized_or), + ErrorSpec{/*aabs=*/0.02, /*arel=*/0.05}, + /*run_hlo_passes=*/false); + status = run_result ? "SUCCESS" : "FAILURE"; + if (!run_result) { + LOG(ERROR) << "RunAndCompare error: " << run_result.message(); + } + } else { + LOG(ERROR) << "ref_optimized error: " << ref_optimized_or.status(); + } + } + } + } + + CuDnnSupportResult actual = CuDnnSupportResult::kNotSupported; + if (optimized_hlo == "__cudnn$fusion" || + optimized_hlo == "__cudnn$blockScaledDot") { + actual = (cudnn_graph_pattern == "BLOCK_SCALE_DEQUANTIZE + MATMUL" && + status == "SUCCESS") + ? CuDnnSupportResult::kSupported + : CuDnnSupportResult::kFailed; + } + + LOG(ERROR) << "Report CuDnnGraph: " << CuDnnSupportResultToString(actual); + LOG(ERROR) << "Report CuDnnGraphPattern: " << cudnn_graph_pattern; + LOG(ERROR) << "Report Status: " << status; + + CuDnnSupportResult expected = param.expected_result; + EXPECT_EQ(actual, expected); + + switch (expected) { + case CuDnnSupportResult::kSupported: + EXPECT_EQ(optimized_hlo, "__cudnn$fusion"); + EXPECT_EQ(cudnn_graph_pattern, "BLOCK_SCALE_DEQUANTIZE + MATMUL"); + EXPECT_EQ(status, "SUCCESS"); + break; + case CuDnnSupportResult::kFailed: + EXPECT_EQ(optimized_hlo, "__cudnn$fusion"); + EXPECT_EQ(cudnn_graph_pattern, "BLOCK_SCALE_DEQUANTIZE + MATMUL"); + EXPECT_EQ(status, "FAILURE"); + break; + case CuDnnSupportResult::kNotSupported: + EXPECT_EQ(optimized_hlo, "dot"); + EXPECT_EQ(cudnn_graph_pattern, "N/A"); + EXPECT_EQ(status, "SUCCESS"); + break; + } +} + +INSTANTIATE_TEST_SUITE_P(CudnnScaledDotCoverageTestSuite, + CudnnScaledDotCoverageTest, + ::testing::ValuesIn(GetCoverageTestCases()), + ScaledDotCoverageTestParamToString); + +} // namespace +} // namespace gpu +} // namespace xla diff --git a/third_party/xla/xla/backends/gpu/codegen/tools/BUILD b/third_party/xla/xla/backends/gpu/codegen/tools/BUILD index b14db5fafbf8f4..51e8db082d5405 100644 --- a/third_party/xla/xla/backends/gpu/codegen/tools/BUILD +++ b/third_party/xla/xla/backends/gpu/codegen/tools/BUILD @@ -108,9 +108,9 @@ xla_cc_binary( "@llvm-project//mlir:Pass", "@llvm-project//mlir:Support", "@llvm-project//mlir:TranslateLib", + "@tensor_ir//:NVTensorIRConversionOptions", "@tensor_ir//:NVTensorIRCudaTilePipelines", "@tensor_ir//:NVTensorIRDialect", - "@tensor_ir//:NVTensorIRToCudaTileConversion", ], ) diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/support.cc b/third_party/xla/xla/backends/gpu/codegen/triton/support.cc index 42d2deeca92700..c9d4e9d9d3a6ad 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/support.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/support.cc @@ -937,6 +937,7 @@ bool IsTritonUnsupportedOpcode(HloOpcode opcode) { case HloOpcode::kScatter: case HloOpcode::kSelectAndScatter: case HloOpcode::kSetDimensionSize: + case HloOpcode::kShuffle: return true; default: return false; diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/support_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/support_test.cc index e30edaf4702c5f..8ee5570c30a039 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/support_test.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/support_test.cc @@ -3737,6 +3737,7 @@ constexpr std::array kUnsupportedOps = { HloOpcode::kScatter, HloOpcode::kSelectAndScatter, HloOpcode::kSetDimensionSize, + HloOpcode::kShuffle, // go/keep-sorted end // clang-format on }; diff --git a/third_party/xla/xla/backends/gpu/collectives/nccl_symmetric_memory.cc b/third_party/xla/xla/backends/gpu/collectives/nccl_symmetric_memory.cc index 49c3433a4f76b9..6531960fb4d149 100644 --- a/third_party/xla/xla/backends/gpu/collectives/nccl_symmetric_memory.cc +++ b/third_party/xla/xla/backends/gpu/collectives/nccl_symmetric_memory.cc @@ -35,7 +35,7 @@ limitations under the License. // Include NCCL after XLA headers. #include "third_party/nccl/nccl.h" -#include "third_party/nccl/nccl_device.h" +#include "third_party/nccl/nccl_device.h" // IWYU pragma: keep #include "xla/tsl/concurrency/executor.h" namespace xla::gpu { @@ -79,16 +79,23 @@ NcclSymmetricMemory::Create(std::shared_ptr comm_state, stream_executor::DeviceAddressBase addr, const std::shared_ptr executor, stream_executor::StreamExecutor* stream_executor) { - ncclWindow_t win; + if (comm_state == nullptr) { + return absl::InvalidArgumentError("NCCL comm_state is null."); + } + + ncclWindow_t win = nullptr; ncclResult_t nccl_status; { - VLOG(3) << absl::StrFormat( - "Create NCCL symmetric memory on comm=%p from: ptr=%p; size=%ld", - comm_state->comm, addr.opaque(), addr.size()); absl::MutexLock lock(comm_state->mutex); - nccl_status = - ncclCommWindowRegister(comm_state->comm, addr.opaque(), addr.size(), - &win, NCCL_WIN_COLL_SYMMETRIC); + ncclComm_t comm = comm_state->comm; + if (comm == nullptr) { + return absl::InvalidArgumentError("NCCL comm is null."); + } + VLOG(3) << absl::StrFormat( + "Create NCCL symmetric memory on comm=%p from: ptr=%p; size=%ld", comm, + addr.opaque(), addr.size()); + nccl_status = ncclCommWindowRegister(comm, addr.opaque(), addr.size(), &win, + NCCL_WIN_COLL_SYMMETRIC); } if (nccl_status != ncclSuccess) { @@ -96,11 +103,22 @@ NcclSymmetricMemory::Create(std::shared_ptr comm_state, return XLA_NCCL_STATUS(nccl_status); } + if (win == nullptr) { + LogInterconnectStatus(stream_executor); + return absl::FailedPreconditionError( + "ncclCommWindowRegister returned a null window because NCCL symmetric " + "memory is not supported on this communicator (e.g. P2P/NVLink or " + "CUMEM is unavailable on fractional/virtualized GPU slices)."); + } + return absl::WrapUnique( new NcclSymmetricMemory(comm_state, win, addr, executor)); } NcclSymmetricMemory::~NcclSymmetricMemory() { + if (win_ == nullptr) { + return; + } absl::Status status = Execute( [&] { @@ -126,9 +144,13 @@ stream_executor::DeviceAddressBase NcclSymmetricMemory::addr() const { absl::StatusOr NcclSymmetricMemory::multimem_addr() const { #if (NCCL_VERSION_CODE >= 22900) || defined(USE_NCCL_HOST_API) + if (win_ == nullptr) { + return absl::FailedPreconditionError( + "Cannot get multimem address from null NCCL window."); + } void* multimem = nullptr; XLA_NCCL_RETURN_IF_ERROR(ncclGetLsaMultimemDevicePointer(win_, 0, &multimem)); - if (multimem) { + if (multimem != nullptr) { return stream_executor::DeviceAddressBase(multimem, addr_.size()); } #endif @@ -139,10 +161,14 @@ NcclSymmetricMemory::multimem_addr() const { absl::StatusOr NcclSymmetricMemory::peer_addr(RankId peer) const { #if (NCCL_VERSION_CODE >= 22902) || defined(USE_NCCL_HOST_API) + if (win_ == nullptr) { + return absl::FailedPreconditionError( + "Cannot get peer address from null NCCL window."); + } void* peer_addr = nullptr; XLA_NCCL_RETURN_IF_ERROR( ncclGetPeerDevicePointer(win_, 0, peer.value(), &peer_addr)); - if (peer_addr) { + if (peer_addr != nullptr) { return stream_executor::DeviceAddressBase(peer_addr, addr_.size()); } return absl::FailedPreconditionError(absl::StrFormat( diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD index d96b6dfca36538..789e090da66f79 100644 --- a/third_party/xla/xla/backends/gpu/runtime/BUILD +++ b/third_party/xla/xla/backends/gpu/runtime/BUILD @@ -1,6 +1,8 @@ load("@local_config_cuda//cuda:build_defs.bzl", "cuda_library") +load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm_is_configured") load("@rules_cc//cc:cc_library.bzl", "cc_library") load("//xla:xla.default.bzl", "xla_cc_test", "xla_internal", "xla_py_proto_library") +load("//xla/stream_executor:build_defs.bzl", "if_cuda_or_rocm_is_configured") load("//xla/tests:build_defs.bzl", "xla_test") load("//xla/tsl:tsl.bzl", "internal_visibility", "nvtx_headers") load("//xla/tsl:tsl.default.bzl", "get_compatible_with_portable") @@ -4283,8 +4285,8 @@ xla_test( "xla_amdgpu_any", "no_oss", ], - "gb200": ["broken"], # TODO(b/505963304): re-enable when fixed. - "gb300": ["broken"], # TODO(b/505963304): re-enable when fixed. + "gb200": ["broken"], # TODO(b/505963049): re-enable when fixed. + "gb300": ["broken"], # TODO(b/505963049): re-enable when fixed. }, backends = ["gpu"], disabled_backends = [], @@ -5539,6 +5541,7 @@ cc_library( "//xla/stream_executor:stream_executor_h", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/debugging:symbolize", + "@com_google_absl//absl/functional:overload", "@com_google_absl//absl/hash", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", @@ -5554,7 +5557,14 @@ cc_library( xla_test( name = "record_ffi_test", - srcs = ["record_ffi_test.cc"], + srcs = ["record_ffi_test.cc"] + if_cuda_is_configured( + ["record_ffi_test_cuda.cc"], + ) + if_rocm_is_configured( + ["record_ffi_test_rocm.cc"], + ) + if_cuda_or_rocm_is_configured( + [], + ["record_ffi_test_default.cc"], + ), backends = ["gpu"], deps = [ ":record_ffi", @@ -5566,9 +5576,9 @@ xla_test( "//xla/ffi:invoke", "//xla/ffi:record_ffi", "//xla/ffi/api:c_api", + "//xla/stream_executor:activate_context", "//xla/stream_executor:command_buffer", "//xla/stream_executor:device_address", - "//xla/stream_executor:kernel_spec", "//xla/stream_executor:platform", "//xla/stream_executor:platform_manager", "//xla/stream_executor:stream", @@ -5576,12 +5586,21 @@ xla_test( "//xla/stream_executor/gpu:gpu_init", "//xla/stream_executor/gpu:gpu_test_kernels_fatbin", "//xla/tsl/platform:test", + "@com_google_absl//absl/base", "@com_google_absl//absl/base:no_destructor", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", + "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings:string_view", "@com_google_absl//absl/types:span", "@com_google_googletest//:gtest_main", - ], + ] + if_cuda_is_configured([ + "@local_config_cuda//cuda:cuda_headers", + "//xla/stream_executor/cuda:cuda_status", + ]) + if_rocm_is_configured([ + "@local_config_rocm//rocm:hip", # buildcleaner: keep + "@local_config_rocm//rocm:rocm_headers", + "//xla/stream_executor/rocm:rocm_status", + ]), ) diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/collective_thunk.cc index f4b30ced48c070..0ad331a60d8761 100644 --- a/third_party/xla/xla/backends/gpu/runtime/collective_thunk.cc +++ b/third_party/xla/xla/backends/gpu/runtime/collective_thunk.cc @@ -47,6 +47,7 @@ limitations under the License. #include "xla/hlo/ir/collective_op_group_mode.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_opcode.h" +#include "xla/layout.h" #include "xla/primitive_util.h" #include "xla/runtime/buffer_use.h" #include "xla/runtime/device_id.h" @@ -71,8 +72,6 @@ limitations under the License. namespace xla::gpu { namespace { -static constexpr int64_t kCollectiveMemorySpaceColor = 1; - bool IsTypeSupportedBy(PrimitiveType element_type, Thunk::Kind reduction_op) { switch (element_type) { case S8: @@ -404,13 +403,13 @@ absl::Status CollectiveThunk::Prepare(const PrepareParams& params) { if (CanUseSymmetricBuffer() && config().use_symmetric_buffer) { for (const Buffer& buffer : buffers_) { - if (buffer.source_memory_space == kCollectiveMemorySpaceColor) { + if (buffer.source_memory_space == Layout::kCollectiveMemorySpace) { ABSL_RETURN_IF_ERROR( params.collective_memory_requests->RequestSymmetricAllocation( clique_key, buffer.source_buffer.slice.index())); } - if (buffer.destination_memory_space == kCollectiveMemorySpaceColor) { + if (buffer.destination_memory_space == Layout::kCollectiveMemorySpace) { ABSL_RETURN_IF_ERROR( params.collective_memory_requests->RequestSymmetricAllocation( clique_key, buffer.destination_buffer.slice.index())); diff --git a/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.cc index 2a2bc8f7f20a3b..0ecd9b0d940d1e 100644 --- a/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.cc +++ b/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.cc @@ -91,7 +91,9 @@ std::string GetSymbolName(const void* ptr) { } struct CustomCallRecordState : public CommandState { - std::vector commands; + absl::flat_hash_map> + commands; }; // A per-execution state that holds state for prepare and initialize stages. @@ -606,13 +608,14 @@ absl::StatusOr CustomCallThunk::Record( const bool is_record_create = std::holds_alternative(record_action); - const bool is_record_update = - std::holds_alternative(record_action); - if (is_record_update) { // Copy over commands from state to inline storage. - TF_RET_CHECK(state->commands.size() <= kMaxCommands) - << "Too many commands to fit in inline storage"; - std::copy(state->commands.begin(), state->commands.end(), commands_storage); - num_commands = state->commands.size(); + if (const auto* record_update = std::get_if(&record_action)) { + if (auto it = state->commands.find(record_update->command); + it != state->commands.end()) { + TF_RET_CHECK(it->second.size() <= kMaxCommands) + << "Too many commands to fit in inline storage"; + std::copy(it->second.begin(), it->second.end(), commands_storage); + num_commands = it->second.size(); + } } XLA_FFI_RecordAction action_to_pass = is_record_create @@ -679,19 +682,19 @@ absl::StatusOr CustomCallThunk::Record( command_buffer); } - // Save newly recorded commands to state if this is the Create action - // Must be done after returning from the FFI handler. - if (is_record_create) { - state->commands.assign(commands_storage, commands_storage + num_commands); - } - // Return the last command in the chain for dependency tracking. // If more than one command was recorded, and they are independent, a dummy // node must be added to the command graph by the FFI client so that XLA // can track a single dependency for the entire chain. if (num_commands > 0 && commands_storage[num_commands - 1] != nullptr) { - return reinterpret_cast( - commands_storage[num_commands - 1]); + const auto* sink_command = + reinterpret_cast( + commands_storage[num_commands - 1]); + if (is_record_create) { + state->commands[sink_command].assign(commands_storage, + commands_storage + num_commands); + } + return sink_command; } // No commands were recorded. return nullptr; diff --git a/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk_test.cc index 63b6b6a03e9daf..459c6b35f04648 100644 --- a/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk_test.cc @@ -1175,5 +1175,95 @@ TEST(CustomCallThunkTest, RecordCommandBufferFfiRecordWithEmptyCommand) { EXPECT_EQ(host_update, 90); } +TEST(CustomCallThunkTest, RecordCommandBufferMultipleRecordsInSameBuffer) { + ASSERT_OK_AND_ASSIGN(se::StreamExecutor * executor, GpuExecutor()); + if (executor->GetDeviceDescription().gpu_compute_capability().IsRocm()) { + GTEST_SKIP() << "AddI32 PTX kernel not supported on ROCm."; + } + + CustomCallThunk::OwnedHandlerBundle bundle; + bundle.execute = + ffi::Ffi::BindExecute().To([]() { return absl::OkStatus(); }); + bundle.record = ffi::Ffi::BindRecord() + .Ctx>() + .Arg() + .Arg() + .Ret() + .To(AddI32FfiHandler); + + ASSERT_OK_AND_ASSIGN(auto setup, FfiRecordTestSetup::Create( + executor, std::move(bundle), + "add_i32_ffi_unroll", {10, 20}, {0})); + + RecordTestAlloc alloc_iter1_create(executor); + ASSERT_OK_AND_ASSIGN( + auto slices_iter1_create, + AllocateAndCopy(*setup->stream, alloc_iter1_create, {1, 2}, {0})); + Thunk::ExecuteParams execute_params_iter1_create = + Thunk::ExecuteParams::Create( + ServiceExecutableRunOptions(), *alloc_iter1_create.buffer_allocations, + setup->stream.get(), setup->stream.get(), nullptr, nullptr, nullptr); + + ASSERT_OK_AND_ASSIGN(auto cb, executor->CreateCommandBuffer( + se::CommandBuffer::Mode::kPrimary)); + + // Record two iterations of the same CustomCallThunk into the same command + // buffer using the same CommandStateManager (as unrolled WhileThunk does). + ASSERT_OK_AND_ASSIGN( + const se::CommandBuffer::Command* cmd_0, + setup->thunk->Record(*setup->execute_params, *setup->record_params, + Command::RecordCreate{/*dependencies=*/{}}, + cb.get())); + ASSERT_OK_AND_ASSIGN( + const se::CommandBuffer::Command* cmd_1, + setup->thunk->Record(execute_params_iter1_create, *setup->record_params, + Command::RecordCreate{/*dependencies=*/{cmd_0}}, + cb.get())); + ASSERT_NE(cmd_0, cmd_1); + + ASSERT_OK(cb->Finalize()); + ASSERT_OK(cb->Submit(setup->stream.get())); + ASSERT_OK(setup->stream->BlockHostUntilDone()); + + // Update both recorded nodes with new buffer allocations. + ASSERT_OK(cb->Update()); + RecordTestAlloc alloc_iter0_update(executor); + ASSERT_OK_AND_ASSIGN( + auto slices_iter0_update, + AllocateAndCopy(*setup->stream, alloc_iter0_update, {40, 50}, {0})); + Thunk::ExecuteParams execute_params_iter0_update = + Thunk::ExecuteParams::Create( + ServiceExecutableRunOptions(), *alloc_iter0_update.buffer_allocations, + setup->stream.get(), setup->stream.get(), nullptr, nullptr, nullptr); + + RecordTestAlloc alloc_iter1_update(executor); + ASSERT_OK_AND_ASSIGN( + auto slices_iter1_update, + AllocateAndCopy(*setup->stream, alloc_iter1_update, {100, 200}, {0})); + Thunk::ExecuteParams execute_params_iter1_update = + Thunk::ExecuteParams::Create( + ServiceExecutableRunOptions(), *alloc_iter1_update.buffer_allocations, + setup->stream.get(), setup->stream.get(), nullptr, nullptr, nullptr); + + ASSERT_OK(setup->thunk->Record(execute_params_iter0_update, + *setup->record_params, + Command::RecordUpdate{cmd_0}, cb.get())); + ASSERT_OK(setup->thunk->Record(execute_params_iter1_update, + *setup->record_params, + Command::RecordUpdate{cmd_1}, cb.get())); + ASSERT_OK(cb->Finalize()); + ASSERT_OK(cb->Submit(setup->stream.get())); + ASSERT_OK(setup->stream->BlockHostUntilDone()); + + int32_t out_0 = 0; + int32_t out_1 = 0; + ASSERT_OK(setup->stream->Memcpy(&out_0, alloc_iter0_update.result_dev_ptrs[0], + sizeof(int32_t))); + ASSERT_OK(setup->stream->Memcpy(&out_1, alloc_iter1_update.result_dev_ptrs[0], + sizeof(int32_t))); + EXPECT_EQ(out_0, 90); + EXPECT_EQ(out_1, 300); +} + } // namespace } // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/runtime/record_ffi.cc b/third_party/xla/xla/backends/gpu/runtime/record_ffi.cc index ef8bf056b05b7e..a423b0881be065 100644 --- a/third_party/xla/xla/backends/gpu/runtime/record_ffi.cc +++ b/third_party/xla/xla/backends/gpu/runtime/record_ffi.cc @@ -24,10 +24,12 @@ limitations under the License. #include #include #include +#include #include #include "absl/container/flat_hash_map.h" #include "absl/debugging/symbolize.h" +#include "absl/functional/overload.h" #include "absl/hash/hash.h" #include "absl/log/check.h" #include "absl/log/log.h" @@ -53,6 +55,7 @@ limitations under the License. #include "tsl/platform/mem.h" namespace xla::gpu { + namespace { template @@ -209,18 +212,24 @@ struct FfiLaunchParams { LaunchDimensions launch_dimensions; std::optional cluster_dims; - se::Kernel* kernel; + std::variant kernel; uint32_t shared_mem_bytes; }; template void AbslStringify(Sink& sink, const FfiLaunchParams& params) { + std::string kernel_name = + std::visit(absl::Overload{[](se::Kernel* k) { return GetSymbolName(k); }, + [](const se::NativeKernel& nk) { + return GetSymbolName(nk.device_fn); + }}, + params.kernel); absl::Format(&sink, "FfiLaunchParams(launch_dimensions=%s, cluster_dims=%s, " "kernel=%s, shared_mem_bytes=%u)", params.launch_dimensions.ToString(), params.cluster_dims.value_or(se::ClusterDim{0, 0, 0}).ToString(), - GetSymbolName(params.kernel), params.shared_mem_bytes); + kernel_name, params.shared_mem_bytes); } class FfiKernelCache : public se::CommandBuffer::Resource { @@ -268,22 +277,35 @@ class FfiKernelCache : public se::CommandBuffer::Resource { << ", ptr: " << it->second.get(); return it->second.get(); } - bool is_ptx = (format == XLA_FFI_SourceFormat_PTX); std::string kernel_name(kernel_name_view); - se::KernelLoaderSpec spec = - is_ptx - ? se::KernelLoaderSpec::CreateCudaPtxInMemorySpec( - AsStringView(kernel_data, kernel_size), kernel_name, num_args) - : se::KernelLoaderSpec::CreateCudaCubinInMemorySpec( - AsByteSpan(kernel_data, kernel_size), kernel_name, num_args); - + std::optional spec; + switch (format) { + case XLA_FFI_SourceFormat_FUNCTION_PTR: { + return absl::InvalidArgumentError( + "XLA_FFI_SourceFormat_FUNCTION_PTR does not use FfiKernelCache"); + } + case XLA_FFI_SourceFormat_PTX: { + spec = se::KernelLoaderSpec::CreateCudaPtxInMemorySpec( + AsStringView(kernel_data, kernel_size), kernel_name, num_args); + break; + } + case XLA_FFI_SourceFormat_CUBIN: { + spec = se::KernelLoaderSpec::CreateCudaCubinInMemorySpec( + AsByteSpan(kernel_data, kernel_size), kernel_name, num_args); + break; + } + } + if (!spec.has_value()) { + return absl::InvalidArgumentError( + absl::StrFormat("Unsupported XLA_FFI_SourceFormat: %d", format)); + } ABSL_ASSIGN_OR_RETURN(std::unique_ptr kernel, - executor->LoadKernel(spec)); - - se::Kernel* kernel_ptr = kernel.get(); - kernels_[KernelKey{kernel_name, kernel_data}] = std::move(kernel); + executor->LoadKernel(*spec)); + auto [inserted_it, unused] = kernels_.try_emplace( + KernelKey{std::move(kernel_name), kernel_data}, std::move(kernel)); + se::Kernel* kernel_ptr = inserted_it->second.get(); XLA_VLOG_DEVICE(3, device_ordinal) - << "FfiKernelCache: created kernel: " << kernel_name + << "FfiKernelCache: created kernel: " << kernel_name_view << ", ptr: " << kernel_ptr; return kernel_ptr; } @@ -326,16 +348,31 @@ XLA_FFI_Error* FfiCreateLaunch( const XLA_FFI_KernelArgs* args, const XLA_FFI_Command* const* dependencies, uint32_t num_dependencies, const XLA_FFI_Command** out_command) { se::CommandBuffer* cmd_buffer = ctx->command_buffer; + const bool pdl = ctx->pdl_enabled && (use_pdl != 0); auto* cache = cmd_buffer->GetOrConstructResource(); - auto kernel_or = - cache->GetOrCreateKernel(ctx->executor, kernel_name, kernel_data, - kernel_size, format, args->num_args); - if (!kernel_or.ok()) { - return xla::ffi::CreateError(kernel_or.status()); + std::variant kernel_target; + if (format == XLA_FFI_SourceFormat_FUNCTION_PTR) { + if (kernel_data == nullptr) { + return xla::ffi::CreateError(absl::InvalidArgumentError( + "kernel_data cannot be null for XLA_FFI_SourceFormat_FUNCTION_PTR")); + } + // FFI C API passes kernel_data as const void*, whereas driver function + // handles (e.g. CUfunction / hipFunction_t) are non-const opaque pointers. + // NOLINTNEXTLINE + kernel_target = se::NativeKernel{const_cast(kernel_data), + std::string(kernel_name), pdl}; + } else { + auto kernel_or = + cache->GetOrCreateKernel(ctx->executor, kernel_name, kernel_data, + kernel_size, format, args->num_args); + if (!kernel_or.ok()) { + return xla::ffi::CreateError(kernel_or.status()); + } + se::Kernel* kernel = *kernel_or; + kernel->set_use_pdl(pdl); + kernel_target = kernel; } - se::Kernel* kernel = *kernel_or; - kernel->set_use_pdl(ctx->pdl_enabled && (use_pdl != 0)); FfiKernelArgsPacked packed_args(args->num_args); packed_args.add_shared_bytes(shared_mem_bytes); @@ -366,19 +403,27 @@ XLA_FFI_Error* FfiCreateLaunch( FfiLaunchParams params{ /*.launch_dimensions = */ LaunchDimensions(blocks, threads), /*.cluster_dims =*/cluster_dims, - /*.kernel =*/kernel, + /*.kernel =*/kernel_target, /*.shared_mem_bytes = */ shared_mem_bytes, }; const int32_t device_ordinal = ctx->executor->device_ordinal(); XLA_VLOG_DEVICE(3, device_ordinal) - << "FfiCreateLaunch for kernel: " << kernel_name - << ", use_pdl: " << kernel->use_pdl() + << "FfiCreateLaunch for kernel: " << kernel_name << ", use_pdl: " << pdl << ", num_dependencies passed: " << num_dependencies << ", resolved deps size: " << deps.size() << ", launch_params: " << params; - auto status_or_cmd = cmd_buffer->CreateLaunch(threads, blocks, cluster_dims, - *kernel, packed_args, deps); + auto status_or_cmd = std::visit( + absl::Overload{ + [&](se::Kernel* k) { + return cmd_buffer->CreateLaunch(threads, blocks, cluster_dims, *k, + packed_args, deps); + }, + [&](const se::NativeKernel& nk) { + return cmd_buffer->CreateLaunch(threads, blocks, cluster_dims, nk, + packed_args, deps); + }}, + params.kernel); if (!status_or_cmd.ok()) { return xla::ffi::CreateError(status_or_cmd.status()); @@ -410,11 +455,22 @@ XLA_FFI_Error* FfiUpdateLaunch(XLA_FFI_RecordContext* ctx, if (absl::Status status = PackArgs(args, packed_args); !status.ok()) { return xla::ffi::CreateError(status); } - if (absl::Status status = cmd_buffer->UpdateLaunch( - cmd, params->launch_dimensions.thread_counts_per_block(), - params->launch_dimensions.block_counts(), params->cluster_dims, - *params->kernel, packed_args); - !status.ok()) { + absl::Status status = std::visit( + absl::Overload{ + [&](se::Kernel* k) { + return cmd_buffer->UpdateLaunch( + cmd, params->launch_dimensions.thread_counts_per_block(), + params->launch_dimensions.block_counts(), params->cluster_dims, + *k, packed_args); + }, + [&](const se::NativeKernel& nk) { + return cmd_buffer->UpdateLaunch( + cmd, params->launch_dimensions.thread_counts_per_block(), + params->launch_dimensions.block_counts(), params->cluster_dims, + nk, packed_args); + }}, + params->kernel); + if (!status.ok()) { return xla::ffi::CreateError(status); } diff --git a/third_party/xla/xla/backends/gpu/runtime/record_ffi_test.cc b/third_party/xla/xla/backends/gpu/runtime/record_ffi_test.cc index eca91c3fc50cde..caebad7b3c0747 100644 --- a/third_party/xla/xla/backends/gpu/runtime/record_ffi_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/record_ffi_test.cc @@ -26,6 +26,7 @@ limitations under the License. #include "absl/log/check.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" +#include "absl/status/statusor.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" #include "xla/ffi/api/record_api.h" @@ -39,7 +40,6 @@ limitations under the License. #include "xla/stream_executor/device_address.h" #include "xla/stream_executor/gpu/gpu_init.h" #include "xla/stream_executor/gpu/gpu_test_kernels_fatbin.h" -#include "xla/stream_executor/kernel_spec.h" #include "xla/stream_executor/platform.h" #include "xla/stream_executor/platform_manager.h" #include "xla/stream_executor/stream.h" @@ -48,6 +48,11 @@ limitations under the License. #include "xla/xla_data.pb.h" namespace xla::gpu { + +absl::StatusOr LoadNativeFunctionPtr(se::StreamExecutor* executor, + absl::Span binary, + absl::string_view name); + namespace { struct KernelBinary { @@ -132,17 +137,7 @@ struct DeviceMemoryBundle { } }; -TEST(RecordFfiTest, KernelLaunchBoundFfi) { - ASSERT_OK_AND_ASSIGN(auto platform, - stream_executor::PlatformManager::PlatformWithName( - stream_executor::GpuPlatformName())); - ASSERT_OK_AND_ASSIGN(auto* executor, platform->ExecutorForDevice(0)); - const auto* cuda_cc = executor->GetDeviceDescription() - .gpu_compute_capability() - .cuda_compute_capability(); - if (cuda_cc && !cuda_cc->IsAtLeastAmpere()) { - GTEST_SKIP() << "Skipping test for compute capability less than Ampere."; - } +void RunRecordFfiTest(se::StreamExecutor* executor, ffi::Ffi& handler) { ASSERT_OK_AND_ASSIGN(auto stream, executor->CreateStream()); std::vector a = {1, 2, 3, 4, 5, 6, 7, 8}; @@ -179,15 +174,7 @@ TEST(RecordFfiTest, KernelLaunchBoundFfi) { ffi::CallFrame call_frame = builder.Build(); - std::unique_ptr handler = - ffi::Ffi::BindRecord() - .Ctx>() - .Arg() - .Arg() - .Arg() - .Arg() - .To(RecordFfiHandler); - ASSERT_OK(ffi::Invoke(ffi::GetXlaFfiApi(), *handler, call_frame, + ASSERT_OK(ffi::Invoke(ffi::GetXlaFfiApi(), handler, call_frame, invoke_context, ffi::ExecutionStage::kRecord)); ASSERT_OK(cmd_buffer->Finalize()); ASSERT_OK(cmd_buffer->Submit(stream.get())); @@ -219,7 +206,7 @@ TEST(RecordFfiTest, KernelLaunchBoundFfi) { ffi::InvokeContext update_context; update_context.extension_start = &record_extension.extension_base; ASSERT_OK(cmd_buffer->Update()); // Begin command buffer update. - ASSERT_OK(ffi::Invoke(ffi::GetXlaFfiApi(), *handler, call_frame, + ASSERT_OK(ffi::Invoke(ffi::GetXlaFfiApi(), handler, call_frame, update_context, ffi::ExecutionStage::kRecord)); ASSERT_OK(cmd_buffer->Finalize()); ASSERT_OK(cmd_buffer->Submit(stream.get())); @@ -232,5 +219,97 @@ TEST(RecordFfiTest, KernelLaunchBoundFfi) { EXPECT_EQ(c, expected_update); } +TEST(RecordFfiTest, KernelLaunchBoundFfi) { + ASSERT_OK_AND_ASSIGN(auto platform, + stream_executor::PlatformManager::PlatformWithName( + stream_executor::GpuPlatformName())); + ASSERT_OK_AND_ASSIGN(auto* executor, platform->ExecutorForDevice(0)); + const auto* cuda_cc = executor->GetDeviceDescription() + .gpu_compute_capability() + .cuda_compute_capability(); + if (cuda_cc && !cuda_cc->IsAtLeastAmpere()) { + GTEST_SKIP() << "Skipping test for compute capability less than Ampere."; + } + + std::unique_ptr handler = + ffi::Ffi::BindRecord() + .Ctx>() + .Arg() + .Arg() + .Arg() + .Arg() + .To(RecordFfiHandler); + RunRecordFfiTest(executor, *handler); +} + +TEST(RecordFfiTest, KernelLaunchWithCuFunc) { + ASSERT_OK_AND_ASSIGN(auto platform, + stream_executor::PlatformManager::PlatformWithName( + stream_executor::GpuPlatformName())); + ASSERT_OK_AND_ASSIGN(auto* executor, platform->ExecutorForDevice(0)); + const auto* cuda_cc = executor->GetDeviceDescription() + .gpu_compute_capability() + .cuda_compute_capability(); + if (cuda_cc && !cuda_cc->IsAtLeastAmpere()) { + GTEST_SKIP() << "Skipping test for compute capability less than Ampere."; + } + + ASSERT_OK_AND_ASSIGN(KernelBinary binary, GetKernelSpec()); + absl::StatusOr func_ptr = + LoadNativeFunctionPtr(executor, binary.bytes, "AddI32"); + if (absl::IsUnimplemented(func_ptr.status())) { + GTEST_SKIP() << func_ptr.status().message(); + } + ASSERT_OK(func_ptr.status()); + + std::unique_ptr handler = + ffi::Ffi::BindRecord() + .Ctx>() + .Arg() + .Arg() + .Arg() + .Arg() + .To([func_ptr = *func_ptr]( + ffi::RecordContext record_ctx, ffi::AnyBuffer input_0, + ffi::AnyBuffer input_1, ffi::AnyBuffer buffer_scratch, + ffi::AnyBuffer result) -> absl::Status { + void* in0 = input_0.untyped_data(); + void* in1 = input_1.untyped_data(); + void* scratch = buffer_scratch.untyped_data(); + void* res = result.untyped_data(); + size_t size = 8 * sizeof(int32_t); + + if (record_ctx.action() == ffi::RecordAction::kCreate) { + ABSL_ASSIGN_OR_RETURN(const XLA_FFI_Command* memcpy_d2d, + record_ctx.CreateMemcpyD2D(scratch, in0, size)); + ABSL_RETURN_IF_ERROR( + record_ctx + .CreateLaunch( + "AddI32", func_ptr, /*kernel_size=*/0, + ffi::SourceFormat::kFunctionPtr, + /*launch_dims=*/{{1, 1, 1}, {8, 1, 1}}, + /*shared_mem_bytes=*/1024, + /*uses_pdl=*/false, + std::vector{ + ffi::DevicePointer{scratch}, + ffi::DevicePointer{in1}, ffi::DevicePointer{res}}, + /*dependencies=*/{memcpy_d2d}) + .status()); + } else if (record_ctx.action() == ffi::RecordAction::kUpdate) { + auto cmds = record_ctx.commands(); + TF_RET_CHECK(cmds.size() == 2) << "Expected 2 commands recorded."; + ABSL_RETURN_IF_ERROR( + record_ctx.UpdateMemcpyD2D(cmds[0], scratch, in0, size)); + ABSL_RETURN_IF_ERROR(record_ctx.UpdateLaunch( + cmds[1], + std::vector{ffi::DevicePointer{scratch}, + ffi::DevicePointer{in1}, + ffi::DevicePointer{res}})); + } + return absl::OkStatus(); + }); + RunRecordFfiTest(executor, *handler); +} + } // namespace } // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/runtime/record_ffi_test_cuda.cc b/third_party/xla/xla/backends/gpu/runtime/record_ffi_test_cuda.cc new file mode 100644 index 00000000000000..856ecf040c1ab4 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/runtime/record_ffi_test_cuda.cc @@ -0,0 +1,46 @@ +/* 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 +#include + +#include "absl/status/status_macros.h" +#include "absl/status/statusor.h" +#include "absl/strings/string_view.h" +#include "absl/types/span.h" +#include "third_party/gpus/cuda/include/cuda.h" +#include "xla/stream_executor/activate_context.h" +#include "xla/stream_executor/cuda/cuda_status.h" +#include "xla/stream_executor/stream_executor.h" + +namespace xla::gpu { + +absl::StatusOr LoadNativeFunctionPtr( + stream_executor::StreamExecutor* executor, absl::Span binary, + absl::string_view name) { + std::unique_ptr activation = + executor->Activate(); + CUmodule module = nullptr; + ABSL_RETURN_IF_ERROR(stream_executor::cuda::ToStatus( + cuModuleLoadData(&module, binary.data()))); + CUfunction function = nullptr; + std::string kernel_name(name); + ABSL_RETURN_IF_ERROR(stream_executor::cuda::ToStatus( + cuModuleGetFunction(&function, module, kernel_name.c_str()))); + return function; +} + +} // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/runtime/record_ffi_test_default.cc b/third_party/xla/xla/backends/gpu/runtime/record_ffi_test_default.cc new file mode 100644 index 00000000000000..ec47494a4fa92e --- /dev/null +++ b/third_party/xla/xla/backends/gpu/runtime/record_ffi_test_default.cc @@ -0,0 +1,33 @@ +/* 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/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/string_view.h" +#include "absl/types/span.h" +#include "xla/stream_executor/stream_executor.h" + +namespace xla::gpu { + +absl::StatusOr LoadNativeFunctionPtr(stream_executor::StreamExecutor*, + absl::Span, + absl::string_view) { + return absl::UnimplementedError( + "Native function pointer loading is not supported on this platform."); +} + +} // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/runtime/record_ffi_test_rocm.cc b/third_party/xla/xla/backends/gpu/runtime/record_ffi_test_rocm.cc new file mode 100644 index 00000000000000..df24728c1442ef --- /dev/null +++ b/third_party/xla/xla/backends/gpu/runtime/record_ffi_test_rocm.cc @@ -0,0 +1,46 @@ +/* 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 +#include + +#include "absl/status/status_macros.h" +#include "absl/status/statusor.h" +#include "absl/strings/string_view.h" +#include "absl/types/span.h" +#include "rocm/include/hip/hip_runtime_api.h" +#include "xla/stream_executor/activate_context.h" +#include "xla/stream_executor/rocm/rocm_status.h" // IWYU pragma: keep +#include "xla/stream_executor/stream_executor.h" + +namespace xla::gpu { + +absl::StatusOr LoadNativeFunctionPtr( + stream_executor::StreamExecutor* executor, absl::Span binary, + absl::string_view name) { + std::unique_ptr activation = + executor->Activate(); + hipModule_t module = nullptr; + ABSL_RETURN_IF_ERROR(stream_executor::gpu::ToStatus( + hipModuleLoadData(&module, binary.data()))); + hipFunction_t function = nullptr; + std::string kernel_name(name); + ABSL_RETURN_IF_ERROR(stream_executor::gpu::ToStatus( + hipModuleGetFunction(&function, module, kernel_name.c_str()))); + return function; +} + +} // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/tests/BUILD b/third_party/xla/xla/backends/gpu/tests/BUILD index 7bc17e886098fe..9775a8ccd7691f 100644 --- a/third_party/xla/xla/backends/gpu/tests/BUILD +++ b/third_party/xla/xla/backends/gpu/tests/BUILD @@ -766,7 +766,6 @@ lit_test_suite_for_gpus( "single_instruction.hlo", "reduce_unnested.hlo", "reduction_vectorization_sm_all.hlo", - "vectorization.hlo", "pdl.hlo", ], "mi300": [ @@ -779,7 +778,6 @@ lit_test_suite_for_gpus( "mi350": [ "dot_bf16.hlo", "single_instruction.hlo", - "vectorization.hlo", "pdl.hlo", ], "mi450": [ @@ -1423,8 +1421,8 @@ xla_test( "gpu": [ "multi_gpu", ], - "gb200": ["broken"], # TODO(b/505963304): re-enable when fixed. - "gb300": ["broken"], # TODO(b/505963304): re-enable when fixed. + "gb200": ["broken"], # TODO(b/505963049): re-enable when fixed. + "gb300": ["broken"], # TODO(b/505963049): re-enable when fixed. }, backends = ["gpu"], deps = [ diff --git a/third_party/xla/xla/backends/gpu/tests/collective_ops_e2e_test.cc b/third_party/xla/xla/backends/gpu/tests/collective_ops_e2e_test.cc index 0f5721059b85ef..15fec5e22fed50 100644 --- a/third_party/xla/xla/backends/gpu/tests/collective_ops_e2e_test.cc +++ b/third_party/xla/xla/backends/gpu/tests/collective_ops_e2e_test.cc @@ -48,6 +48,7 @@ limitations under the License. #include "xla/hlo/testlib/pattern_matcher_gmock.h" #include "xla/hlo/testlib/verified_hlo_module.h" #include "xla/hlo/utils/hlo_matchers.h" +#include "xla/layout.h" #include "xla/literal.h" #include "xla/literal_util.h" #include "xla/service/gpu/backend_configs.pb.h" @@ -3406,9 +3407,11 @@ ENTRY main { test_runner().HloModuleFromWrapped(executable.get())); const HloInstruction* ag_start = FindCollectiveStarts(executable_module, HloOpcode::kAllGather).at(0); - // Both ag and its producer should have collective memory space 1 - EXPECT_EQ(ag_start->shape().tuple_shapes()[1].layout().memory_space(), 1); - EXPECT_EQ(ag_start->operand(0)->shape().layout().memory_space(), 1); + // Both ag and its producer should have collective memory space. + EXPECT_EQ(ag_start->shape().tuple_shapes()[1].layout().memory_space(), + Layout::kCollectiveMemorySpace); + EXPECT_EQ(ag_start->operand(0)->shape().layout().memory_space(), + Layout::kCollectiveMemorySpace); } TEST_F(CollectiveOpsTestE2E, @@ -3457,7 +3460,8 @@ ROOT tuple = (bf16[1024,1024]{1,0}, bf16[]) tuple(all-reduce-done, all-reduce-do // space. for (auto ar : all_ar) { EXPECT_EQ(ar->operand(0)->opcode(), HloOpcode::kCopy); - EXPECT_EQ(ar->operand(0)->shape().layout().memory_space(), 1); + EXPECT_EQ(ar->operand(0)->shape().layout().memory_space(), + Layout::kCollectiveMemorySpace); } } diff --git a/third_party/xla/xla/backends/gpu/tests/collective_ops_ffi_test.cc b/third_party/xla/xla/backends/gpu/tests/collective_ops_ffi_test.cc index 7b02691885318b..8cce5f8a4b1236 100644 --- a/third_party/xla/xla/backends/gpu/tests/collective_ops_ffi_test.cc +++ b/third_party/xla/xla/backends/gpu/tests/collective_ops_ffi_test.cc @@ -1227,8 +1227,8 @@ TEST_P(AllReduceTest, DeviceAllReduce) { ENTRY test_computation { id = u32[] replica-id() - in = u32[]{:S(1)} copy(id) - all-reduce = u32[]{:S(1)} custom-call(in), + in = u32[]{:S(7)} copy(id) + all-reduce = u32[]{:S(7)} custom-call(in), custom_call_target="__xla_test_$0_device_all_reduce", api_version=API_VERSION_TYPED_FFI ROOT out = u32[] copy(all-reduce) @@ -1316,7 +1316,7 @@ TEST_P(AllReduceTest, MulticastAllReduce) { ENTRY test_computation { c0 = u32[] constant(1) - in = u32[]{:S(1)} copy(c0) + in = u32[]{:S(7)} copy(c0) all-reduce = u32[] custom-call(in), custom_call_target="__xla_test_$0_multimem_all_reduce", api_version=API_VERSION_TYPED_FFI @@ -1360,7 +1360,7 @@ TEST_P(AllReduceTest, SymMulticastAllReduce) { ENTRY test_computation { c0 = u32[] constant(1) - in = u32[]{:S(1)} copy(c0) + in = u32[]{:S(7)} copy(c0) all-reduce = u32[] custom-call(in), custom_call_target="__xla_test_$0_sym_multimem_all_reduce", api_version=API_VERSION_TYPED_FFI @@ -1405,7 +1405,7 @@ TEST_P(AllReduceTest, SymPeerAllReduce) { ENTRY test_computation { id = u32[] replica-id() - in = u32[]{:S(1)} copy(id) + in = u32[]{:S(7)} copy(id) all-reduce = u32[] custom-call(in), custom_call_target="__xla_test_$0_sym_peer_all_reduce", api_version=API_VERSION_TYPED_FFI @@ -1442,7 +1442,7 @@ INSTANTIATE_TEST_SUITE_P( }); // Same as DeviceAllReduce, but uses frontend_attributes to specify memory -// spaces instead of hardcoded S(1). +// spaces instead of hardcoded S(7). TEST_F(CollectiveOpsTestFFI, DeviceAllReduceWithFrontendAttributes) { if (device_count() < kNumReplicas) { GTEST_SKIP() << "Test requires at least " << kNumReplicas << " devices (" @@ -1462,8 +1462,8 @@ TEST_F(CollectiveOpsTestFFI, DeviceAllReduceWithFrontendAttributes) { custom_call_target="__xla_test_blocked_device_all_reduce", api_version=API_VERSION_TYPED_FFI, frontend_attributes={ - operands_memory_spaces="{0:1}", - results_memory_spaces="{0:1}" + operands_memory_spaces="{0:7}", + results_memory_spaces="{0:7}" } ROOT out = u32[] copy(all-reduce) } diff --git a/third_party/xla/xla/backends/gpu/tests/vectorization.hlo b/third_party/xla/xla/backends/gpu/tests/vectorization.hlo index 56a4747a3195e8..d2cbe6e7058d1c 100644 --- a/third_party/xla/xla/backends/gpu/tests/vectorization.hlo +++ b/third_party/xla/xla/backends/gpu/tests/vectorization.hlo @@ -1,5 +1,5 @@ // RUN: hlo-opt %s --platform=gpu --stage=llvm-before-optimizations --xla_gpu_target_config_filename=%S/../target_config/specs/%{GPU}.txtpb | FileCheck --check-prefixes=CHECK-LLVM-%{GPU} %s -// RUN: hlo-opt %s --platform=gpu --stage=ptx --xla_gpu_target_config_filename=%S/../target_config/specs/%{GPU}.txtpb | FileCheck --check-prefixes=CHECK-PTX-%{GPU} %s +// RUN: %if !IS_ROCM %{ hlo-opt %s --platform=gpu --stage=ptx --xla_gpu_target_config_filename=%S/../target_config/specs/%{GPU}.txtpb | FileCheck --check-prefixes=CHECK-PTX-%{GPU} %s %} // CHECK-LLVM-b200: store <16 x i8> // CHECK-PTX-b200: st.global.v4.b32 @@ -12,6 +12,8 @@ // CHECK-PTX-a100_pcie_80: st.global.b32 // CHECK-LLVM-a6000: store <4 x i8> // CHECK-PTX-a6000: st.global.b32 +// CHECK-LLVM-mi200: store <4 x i8> +// CHECK-LLVM-mi350: store <4 x i8> HloModule m8, is_scheduled=true diff --git a/third_party/xla/xla/backends/gpu/transforms/BUILD b/third_party/xla/xla/backends/gpu/transforms/BUILD index fcb31cc5e44121..2013ccd9e9dc89 100644 --- a/third_party/xla/xla/backends/gpu/transforms/BUILD +++ b/third_party/xla/xla/backends/gpu/transforms/BUILD @@ -763,6 +763,7 @@ cc_library( deps = [ "//xla/codegen:ir_emission_utils", "//xla/hlo/ir:hlo", + "//xla/hlo/ir:hlo_instruction_utils", "//xla/hlo/pass:hlo_pass", "//xla/hlo/utils:hlo_traversal", "//xla/service:call_graph", diff --git a/third_party/xla/xla/backends/gpu/transforms/copy_fusion.cc b/third_party/xla/xla/backends/gpu/transforms/copy_fusion.cc index e3d9e9c2943cb5..072800c0eb340a 100644 --- a/third_party/xla/xla/backends/gpu/transforms/copy_fusion.cc +++ b/third_party/xla/xla/backends/gpu/transforms/copy_fusion.cc @@ -29,6 +29,7 @@ limitations under the License. #include "absl/strings/string_view.h" #include "xla/codegen/ir_emission_utils.h" #include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_instruction_utils.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/ir/hlo_original_value.h" #include "xla/hlo/utils/hlo_traversal.h" @@ -53,7 +54,9 @@ bool OnlyElementwiseOpsReachableFromParams(HloComputation* fused_computation) { for (auto user : hlo->users()) { if ((!user->IsElementwiseOnOperand(user->operand_index(hlo)) || HloPredicateIsOp(user)) && - HloPredicateIsNotOp(user)) { + !(HloPredicateIsOp(user) && + hlo_instruction_utils::KeepsBitwidth(*user)) && + HloPredicateIsNotOp(user)) { return false; } if (visited.insert(user).second) { @@ -101,9 +104,10 @@ absl::StatusOr CopyFusion::DoCopyFusion( } copy_user = copy_user->users()[0]; } - // Skip bitcast ops. + // Skip bitcast ops that keep the bitwidth. if (HloPredicateIsOp(copy_user) && - copy_user->user_count() == 1) { + copy_user->user_count() == 1 && + hlo_instruction_utils::KeepsBitwidth(*copy_user)) { copy_user = copy_user->users()[0]; } if (HloPredicateIsOp(copy_user) && diff --git a/third_party/xla/xla/backends/gpu/transforms/copy_fusion_test.cc b/third_party/xla/xla/backends/gpu/transforms/copy_fusion_test.cc index ddc4950643d6e8..7780daebd6efc9 100644 --- a/third_party/xla/xla/backends/gpu/transforms/copy_fusion_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/copy_fusion_test.cc @@ -676,5 +676,49 @@ TEST_F(CopyFusionTest, PropagateOriginalValue) { EXPECT_EQ(fusion->original_value()->ToString(), R"(({"transpose"}, {}, {}))"); } +TEST_F(CopyFusionTest, CopyFusionBitcastConvertSameBitwidthIsFused) { + ASSERT_OK_AND_ASSIGN( + auto module, ParseAndReturnVerifiedModule(absl::StrCat(kModulePrefix, R"( + fused_computation { + p0 = f32[16,32]{1,0} parameter(0) + ROOT neg = f32[16,32]{1,0} negate(p0) + } + + ENTRY main { + p0 = f32[16,32]{1,0} parameter(0) + fusion = f32[16,32]{1,0} fusion(p0), kind=kLoop, calls=fused_computation + bitcast = s32[16,32]{1,0} bitcast(fusion) + copy = s32[16,32]{1,0} copy(bitcast) + ROOT t = (f32[16,32]{1,0}, s32[16,32]{1,0}) tuple(fusion, copy) + })"))); + ASSERT_OK_AND_ASSIGN(bool changed, cf_.Run(module.get())); + EXPECT_TRUE(changed); + const HloInstruction* root = module->entry_computation()->root_instruction(); + const HloInstruction* fusion = nullptr; + ASSERT_THAT(root, GmockMatch(m::Tuple(m::GetTupleElement(m::Fusion(&fusion)), + m::GetTupleElement()))); + EXPECT_THAT(fusion->fused_expression_root(), + GmockMatch(m::Tuple(m::Negate(), m::Copy(m::Bitcast())))); +} + +TEST_F(CopyFusionTest, CopyFusionBitcastConvertDifferentBitwidthIsNotFused) { + ASSERT_OK_AND_ASSIGN( + auto module, ParseAndReturnVerifiedModule(absl::StrCat(kModulePrefix, R"( + fused_computation { + p0 = f32[16,32]{1,0} parameter(0) + ROOT neg = f32[16,32]{1,0} negate(p0) + } + + ENTRY main { + p0 = f32[16,32]{1,0} parameter(0) + fusion = f32[16,32]{1,0} fusion(p0), kind=kLoop, calls=fused_computation + bitcast = s8[16,32,4]{2,1,0} bitcast(fusion) + copy = s8[16,32,4]{2,1,0} copy(bitcast) + ROOT t = (f32[16,32]{1,0}, s8[16,32,4]{2,1,0}) tuple(fusion, copy) + })"))); + ASSERT_OK_AND_ASSIGN(bool changed, cf_.Run(module.get())); + EXPECT_FALSE(changed); +} + } // namespace gpu } // namespace xla diff --git a/third_party/xla/xla/codegen/emitters/BUILD b/third_party/xla/xla/codegen/emitters/BUILD index b3e3a248ebc56e..0e987b5d9a85fc 100644 --- a/third_party/xla/xla/codegen/emitters/BUILD +++ b/third_party/xla/xla/codegen/emitters/BUILD @@ -224,7 +224,7 @@ cc_library( hdrs = ["kernel_arguments.h"], visibility = [ ":friends", - "//xla/backends/gpu/libraries/native_custom_call_thunks:handler_allowlist", + "//xla:native_custom_call_handler_allowlist", ], deps = [ "//xla:frontend_attributes", diff --git a/third_party/xla/xla/codegen/tiling/experimental/BUILD b/third_party/xla/xla/codegen/tiling/experimental/BUILD index 1c3c1cefb0ab2a..fd9be74a3cf56d 100644 --- a/third_party/xla/xla/codegen/tiling/experimental/BUILD +++ b/third_party/xla/xla/codegen/tiling/experimental/BUILD @@ -198,6 +198,7 @@ xla_cc_test( "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/strings:string_view", "@com_google_googletest//:gtest_main", + "@llvm-project//llvm:Support", "@llvm-project//mlir:IR", ], ) diff --git a/third_party/xla/xla/codegen/tiling/experimental/tile.cc b/third_party/xla/xla/codegen/tiling/experimental/tile.cc index 20947f5d663f00..15623bf136dbd3 100644 --- a/third_party/xla/xla/codegen/tiling/experimental/tile.cc +++ b/third_party/xla/xla/codegen/tiling/experimental/tile.cc @@ -19,6 +19,7 @@ limitations under the License. #include #include +#include "absl/log/check.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/str_cat.h" @@ -225,17 +226,51 @@ void Tile::Replace(const llvm::DenseMap& map) { } } +void SimplifyDimTiles( + llvm::ArrayRef> dim_tile_groups, + const TilingSpace& space) { + int64_t total_dim_tiles = 0; + for (const auto& group : dim_tile_groups) { + total_dim_tiles += group.size(); + } + if (total_dim_tiles == 0) { + return; + } + llvm::SmallVector expressions; + expressions.reserve(total_dim_tiles * 4); + for (auto& group : dim_tile_groups) { + for (DimTile& dim_tile : group) { + expressions.push_back(dim_tile.offset); + expressions.push_back(dim_tile.size); + expressions.push_back(dim_tile.stride); + expressions.push_back(dim_tile.upper_bound); + } + } + expressions = space.SimplifyExpressions(expressions); + CHECK_EQ(expressions.size(), total_dim_tiles * 4); + int idx = 0; + for (auto& group : dim_tile_groups) { + for (DimTile& dim_tile : group) { + dim_tile.offset = expressions[idx++]; + dim_tile.size = expressions[idx++]; + dim_tile.stride = expressions[idx++]; + dim_tile.upper_bound = expressions[idx++]; + } + } +} + +void SimplifyDimTiles(llvm::MutableArrayRef dim_tiles, + const TilingSpace& space) { + SimplifyDimTiles( + llvm::ArrayRef>(&dim_tiles, 1), space); +} + void DimTile::Simplify(const TilingSpace& space) { - offset = space.SimplifyExpression(offset); - size = space.SimplifyExpression(size); - stride = space.SimplifyExpression(stride); - upper_bound = space.SimplifyExpression(upper_bound); + SimplifyDimTiles(*this, space); } void Tile::Simplify() { - for (DimTile& dim_tile : llvm::concat(dim_tiles_, replica_ids_)) { - dim_tile.Simplify(*tiling_space_); - } + SimplifyDimTiles({dim_tiles_, replica_ids_}, *tiling_space_); } Tile Tile::CloneWithNewDims(llvm::SmallVector new_dim_tiles) const { diff --git a/third_party/xla/xla/codegen/tiling/experimental/tile.h b/third_party/xla/xla/codegen/tiling/experimental/tile.h index dbe519ae89df6e..3c958a17d30489 100644 --- a/third_party/xla/xla/codegen/tiling/experimental/tile.h +++ b/third_party/xla/xla/codegen/tiling/experimental/tile.h @@ -143,6 +143,17 @@ H AbslHashValue(H h, const DimTile& dim_tile) { return H::combine(std::move(h), static_cast(dim_tile_hash)); } +// Simplifies groups of DimTiles using the dimension and symbol bounds of the +// given tiling space. +void SimplifyDimTiles( + llvm::ArrayRef> dim_tile_groups, + const TilingSpace& space); + +// Simplifies a list of DimTiles using the dimension and symbol bounds of the +// given tiling space. +void SimplifyDimTiles(llvm::MutableArrayRef dim_tiles, + const TilingSpace& space); + // Tile is a collection of tilings for every dimension of output tensor // of an HLO instruction. TiledHloInstruction associates a Tile // with an HLO instruction. diff --git a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc index edd62cb0501e6d..8b2abc3b0c8219 100644 --- a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc +++ b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc @@ -1340,14 +1340,13 @@ absl::Status PropagateTileThroughMinimalReshape( mlir_context) + 1; - target_dim_tiles[target_id].Simplify(tiling_space); - if (tiling_space.IsSymbolic()) { VLOG(2) << "Skipping reshape contiguity check as tile sizes are not " "assigned yet."; return absl::OkStatus(); } + target_dim_tiles[target_id].Simplify(tiling_space); return VerifyReshapeContiguity( source_info.dims, source_info.tiles, target_info.dims, absl::MakeSpan(&target_dim_tiles[target_id], 1)); @@ -1369,8 +1368,6 @@ absl::Status PropagateTileThroughMinimalReshape( target_dim_tiles[target_id].size = upper_bounds_inclusive[i] - offsets[i] + 1; target_dim_tiles[target_id].upper_bound = upper_bounds_inclusive[i] + 1; - target_dim_tiles[target_id].Simplify(tiling_space); - target_info.tiles[i] = target_dim_tiles[target_id]; } if (tiling_space.IsSymbolic()) { @@ -1379,6 +1376,14 @@ absl::Status PropagateTileThroughMinimalReshape( return absl::OkStatus(); } + for (auto [i, target_id] : llvm::enumerate(target_info.ids)) { + target_info.tiles[i] = target_dim_tiles[target_id]; + } + SimplifyDimTiles(target_info.tiles, tiling_space); + for (auto [i, target_id] : llvm::enumerate(target_info.ids)) { + target_dim_tiles[target_id] = target_info.tiles[i]; + } + return VerifyReshapeContiguity(source_info.dims, source_info.tiles, target_info.dims, target_info.tiles); } diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc index 12f46e69b2e2e7..f010e9db706c48 100644 --- a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc +++ b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc @@ -618,18 +618,24 @@ void TilingSpace::InitSimplificationIndexing() { } } -SymbolicExpr TilingSpace::SimplifyExpression(const SymbolicExpr& expr) const { +llvm::SmallVector TilingSpace::SimplifyExpressions( + const llvm::SmallVector& expressions) const { if (is_symbolic_) { - return expr.Canonicalize(); + llvm::SmallVector simplified_expressions; + simplified_expressions.reserve(expressions.size()); + for (const auto& expr : expressions) { + simplified_expressions.push_back(expr.Canonicalize()); + } + return simplified_expressions; } - + // TODO(b/565301234): add constraints from tiling space? SymbolicMap map = SymbolicMap::Get(mlir_context(), dimensions_.size(), - rt_vars_.size(), {expr}); + rt_vars_.size(), expressions); IndexingMap indexing_map(map, dim_vars_indexing_, range_vars_indexing_, rt_vars_indexing_); indexing_map.Simplify(IndexingMap::SimplifyPointDimensions::kPreserve); - return indexing_map.GetSymbolicMap().GetResults()[0]; + return std::move(indexing_map).GetSymbolicMap().GetResults(); } absl::StatusOr>> diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h index fb7f920c8757cf..4c1c5a24b50234 100644 --- a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h +++ b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h @@ -194,9 +194,10 @@ class TilingSpace { bool IsSymbolic() const { return is_symbolic_; } - // Simplifies an expression using actual dimension and symbol bounds + // Simplifies expressions using actual dimension and symbol bounds // based on the assigned tile sizes and runtime variable bounds. - SymbolicExpr SimplifyExpression(const SymbolicExpr& expr) const; + llvm::SmallVector SimplifyExpressions( + const llvm::SmallVector& expressions) const; // Returns the list of valid tilings for the tiling space. absl::StatusOr>> GetValidTilings(); diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc b/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc index 5ff15bdce99216..8bb35fcda03251 100644 --- a/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc +++ b/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc @@ -26,7 +26,9 @@ limitations under the License. #include "absl/status/status.h" #include "absl/status/status_matchers.h" #include "absl/strings/string_view.h" +#include "llvm/ADT/SmallVector.h" #include "mlir/IR/MLIRContext.h" +#include "xla/codegen/tiling/experimental/tile.h" #include "xla/hlo/analysis/indexing_test_utils.h" #include "xla/hlo/analysis/symbolic_expr.h" #include "xla/hlo/analysis/symbolic_map.h" @@ -42,6 +44,7 @@ namespace { using ::absl_testing::StatusIs; using ::mlir::MLIRContext; +using ::testing::ElementsAre; using ::testing::HasSubstr; MATCHER_P(MatchString, tiling_space_string, "") { @@ -441,7 +444,7 @@ TEST_F(TilingSpaceSameShapeMultiOutputTest, TwoOutputsEqualShapesParallelDims) { )")); } -class TilingSpaceSimplifyExpressionTest : public TilingSpaceTest { +class TilingSpaceSimplifyExpressionsTest : public TilingSpaceTest { public: void SetUp() override { TilingSpaceTest::SetUp(); @@ -463,41 +466,71 @@ class TilingSpaceSimplifyExpressionTest : public TilingSpaceTest { CHECK_OK(tiling_space_->AssignTileSizes({16, 2})); } + SymbolicExpr ParseExpr(absl::string_view expr_str) { + return ParseSymbolicExpr(expr_str, &mlir_context_, /*num_dims=*/2); + } + std::unique_ptr tiling_space_; }; -TEST_F(TilingSpaceSimplifyExpressionTest, ModRemovedIfLessThanDivisor) { - SymbolicExpr tid_0 = CreateDimExpr(0, &mlir_context_); - EXPECT_EQ(tiling_space_->SimplifyExpression((tid_0 * 8) % 96), tid_0 * 8); +TEST_F(TilingSpaceSimplifyExpressionsTest, ModRemovedIfLessThanDivisor) { + EXPECT_THAT( + tiling_space_->SimplifyExpressions({ParseExpr("(d0 * 8) mod 96")}), + ElementsAre(ParseExpr("d0 * 8"))); } -TEST_F(TilingSpaceSimplifyExpressionTest, FloorDivFactorsDivisor) { - SymbolicExpr tid_1 = CreateDimExpr(1, &mlir_context_); - EXPECT_EQ(tiling_space_->SimplifyExpression((tid_1 * 2).floorDiv(10)), - tid_1.floorDiv(5)); +TEST_F(TilingSpaceSimplifyExpressionsTest, MultipleExpressionsSimplified) { + EXPECT_THAT(tiling_space_->SimplifyExpressions({ParseExpr("(d0 * 8) mod 96"), + ParseExpr("(d1 * 2) / 10"), + ParseExpr("d0 * 16 + 500")}), + ElementsAre(ParseExpr("d0 * 8"), ParseExpr("d1 / 5"), + ParseExpr("d0 * 16 + 500"))); } -TEST_F(TilingSpaceSimplifyExpressionTest, - ExpressionUnchangedIfNotAlgebraicallyFolds) { - SymbolicExpr tid_0 = CreateDimExpr(0, &mlir_context_); - EXPECT_EQ(tiling_space_->SimplifyExpression(tid_0 * 16 + 500), - tid_0 * 16 + 500); +TEST_F(TilingSpaceSimplifyExpressionsTest, DimTileSimplify) { + DimTile dt{ParseExpr("(d0 * 8) mod 96"), ParseExpr("(d1 * 2) / 10"), + ParseExpr("d0 * 16 + 500"), + ParseExpr("(d0 * 16 + d1 * 2) mod 200")}; + dt.Simplify(*tiling_space_); + EXPECT_EQ(dt.offset, ParseExpr("d0 * 8")); + EXPECT_EQ(dt.size, ParseExpr("d1 / 5")); + EXPECT_EQ(dt.stride, ParseExpr("d0 * 16 + 500")); + EXPECT_EQ(dt.upper_bound, ParseExpr("d0 * 16 + d1 * 2")); } -TEST_F(TilingSpaceSimplifyExpressionTest, NestedFloorDivFactorsDivisor) { - auto expr = ParseSymbolicExpr("(d0 * 16 + d1 * 2) / 200", &mlir_context_, - /*num_dims=*/2); - EXPECT_EQ(tiling_space_->SimplifyExpression(expr), - ParseSymbolicExpr("(d0 * 8 + d1) / 100", &mlir_context_, - /*num_dims=*/2)); +TEST_F(TilingSpaceSimplifyExpressionsTest, SimplifyDimTiles) { + llvm::SmallVector dim_tiles = { + {ParseExpr("(d0 * 8) mod 96"), ParseExpr("(d1 * 2) / 10"), + ParseExpr("d0 * 16 + 500"), ParseExpr("(d0 * 16 + d1 * 2) mod 200")}, + {ParseExpr("(d1 * 2) / 10"), ParseExpr("(d0 * 8) mod 96"), + ParseExpr("d0 * 16 + 500"), ParseExpr("(d0 * 16 + d1 * 2) mod 200")}}; + SimplifyDimTiles(dim_tiles, *tiling_space_); + EXPECT_EQ(dim_tiles[0].offset, ParseExpr("d0 * 8")); + EXPECT_EQ(dim_tiles[0].size, ParseExpr("d1 / 5")); + EXPECT_EQ(dim_tiles[0].stride, ParseExpr("d0 * 16 + 500")); + EXPECT_EQ(dim_tiles[0].upper_bound, ParseExpr("d0 * 16 + d1 * 2")); + EXPECT_EQ(dim_tiles[1].offset, ParseExpr("d1 / 5")); + EXPECT_EQ(dim_tiles[1].size, ParseExpr("d0 * 8")); + EXPECT_EQ(dim_tiles[1].stride, ParseExpr("d0 * 16 + 500")); + EXPECT_EQ(dim_tiles[1].upper_bound, ParseExpr("d0 * 16 + d1 * 2")); + + llvm::SmallVector empty_dim_tiles; + SimplifyDimTiles(empty_dim_tiles, *tiling_space_); + EXPECT_TRUE(empty_dim_tiles.empty()); } -TEST_F(TilingSpaceSimplifyExpressionTest, NestedModRemovedIfLessThanDivisor) { - auto expr = ParseSymbolicExpr("(d0 * 16 + d1 * 2) mod 200", &mlir_context_, - /*num_dims=*/2); - EXPECT_EQ( - tiling_space_->SimplifyExpression(expr), - ParseSymbolicExpr("d0 * 16 + d1 * 2", &mlir_context_, /*num_dims=*/2)); +TEST_F(TilingSpaceSimplifyExpressionsTest, SimplifyDimTilesGroups) { + llvm::SmallVector group1 = { + {ParseExpr("(d0 * 8) mod 96"), ParseExpr("(d1 * 2) / 10"), + ParseExpr("d0 * 16 + 500"), ParseExpr("(d0 * 16 + d1 * 2) mod 200")}}; + llvm::SmallVector group2 = { + {ParseExpr("(d1 * 2) / 10"), ParseExpr("(d0 * 8) mod 96"), + ParseExpr("d0 * 16 + 500"), ParseExpr("(d0 * 16 + d1 * 2) mod 200")}}; + SimplifyDimTiles({group1, group2}, *tiling_space_); + EXPECT_EQ(group1[0].offset, ParseExpr("d0 * 8")); + EXPECT_EQ(group1[0].size, ParseExpr("d1 / 5")); + EXPECT_EQ(group2[0].offset, ParseExpr("d1 / 5")); + EXPECT_EQ(group2[0].size, ParseExpr("d0 * 8")); } TEST_F(TilingSpaceTest, ClonePerformsDeepCopies) { diff --git a/third_party/xla/xla/codegen/xtile/ir/tests/ops.mlir b/third_party/xla/xla/codegen/xtile/ir/tests/ops.mlir index 4a8457949af424..d05c7e37e3741d 100644 --- a/third_party/xla/xla/codegen/xtile/ir/tests/ops.mlir +++ b/third_party/xla/xla/codegen/xtile/ir/tests/ops.mlir @@ -282,3 +282,36 @@ func.func @memref_layout_shape_size_mismatch(%arg0: memref<1024xf32, #xtile.layo func.func @memref_layout_is_not_a_permutation(%arg0: memref<1024xf32, #xtile.layout<[1]>>) { return } + +// ----- + +func.func @memref_bitcast(%arg0: memref<4x16xbf16, #xtile.layout<[0, 1]>>) + -> memref<4x16xi16, #xtile.layout<[0, 1]>> { + %0 = xtile.memref_bitcast %arg0 : memref<4x16xbf16, #xtile.layout<[0, 1]>> -> memref<4x16xi16, #xtile.layout<[0, 1]>> + return %0 : memref<4x16xi16, #xtile.layout<[0, 1]>> +} + +// ----- + +func.func @memref_bitcast_shape_mismatch(%arg0: memref<4x16xbf16>) -> memref<16x4xi16> { + // expected-error@+1 {{requires the same shape for all operands and results}} + %0 = xtile.memref_bitcast %arg0 : memref<4x16xbf16> -> memref<16x4xi16> + return %0 : memref<16x4xi16> +} + +// ----- + +func.func @memref_bitcast_layout_mismatch(%arg0: memref<4x16xbf16, #xtile.layout<[0, 1]>>) + -> memref<4x16xi16, #xtile.layout<[1, 0]>> { + // expected-error@+1 {{does not match result layout}} + %0 = xtile.memref_bitcast %arg0 : memref<4x16xbf16, #xtile.layout<[0, 1]>> -> memref<4x16xi16, #xtile.layout<[1, 0]>> + return %0 : memref<4x16xi16, #xtile.layout<[1, 0]>> +} + +// ----- + +func.func @memref_bitcast_width_mismatch(%arg0: memref<16xbf16>) -> memref<16xi32> { + // expected-error@+1 {{element bit widths differ: 'bf16' vs 'i32'}} + %0 = xtile.memref_bitcast %arg0 : memref<16xbf16> -> memref<16xi32> + return %0 : memref<16xi32> +} diff --git a/third_party/xla/xla/codegen/xtile/ir/xtile_ops.cc b/third_party/xla/xla/codegen/xtile/ir/xtile_ops.cc index aada99430a4ce4..3e6a2c34a280af 100644 --- a/third_party/xla/xla/codegen/xtile/ir/xtile_ops.cc +++ b/third_party/xla/xla/codegen/xtile/ir/xtile_ops.cc @@ -177,6 +177,50 @@ mlir::LogicalResult EntryFuncOp::verify() { return mlir::success(); } +mlir::LogicalResult MemRefBitcastOp::verify() { + mlir::MemRefType source_type = getSource().getType(); + mlir::MemRefType result_type = getResult().getType(); + + if (source_type.getLayout() != result_type.getLayout()) { + return emitOpError() << "source layout: " << source_type.getLayout() + << " does not match result layout: " + << result_type.getLayout(); + } + if (source_type.getMemorySpace() != result_type.getMemorySpace()) { + return emitOpError() << "source and result memory spaces differ"; + } + + mlir::Type source_element_type = source_type.getElementType(); + mlir::Type result_element_type = result_type.getElementType(); + if (!source_element_type.isIntOrFloat() || + !result_element_type.isIntOrFloat()) { + return emitOpError() << "element types must be integer or float, got: " + << source_element_type << " and " + << result_element_type; + } + if (source_element_type.getIntOrFloatBitWidth() != + result_element_type.getIntOrFloatBitWidth()) { + return emitOpError() << "element bit widths differ: " << source_element_type + << " vs " << result_element_type; + } + return mlir::success(); +} + +mlir::OpFoldResult MemRefBitcastOp::fold(FoldAdaptor) { + if (getSource().getType() == getType()) { + return getSource(); + } + // bitcast(bitcast(x)) -> x, or a single bitcast from x otherwise. + if (auto producer = getSource().getDefiningOp()) { + if (producer.getSource().getType() == getType()) { + return producer.getSource(); + } + getSourceMutable().assign(producer.getSource()); + return getResult(); + } + return {}; +} + mlir::TypedValue ExtractTileOp::getBuffer() { return getSource(); } diff --git a/third_party/xla/xla/codegen/xtile/ir/xtile_ops.td b/third_party/xla/xla/codegen/xtile/ir/xtile_ops.td index aa03b558989d4a..1cbaab2442e73b 100644 --- a/third_party/xla/xla/codegen/xtile/ir/xtile_ops.td +++ b/third_party/xla/xla/codegen/xtile/ir/xtile_ops.td @@ -187,6 +187,31 @@ def SelectBufferOp : XTile_Op<"select_buffer", [Pure]> { }]; } +def MemRefBitcastOp : XTile_Op<"memref_bitcast", + [Pure, SameOperandsAndResultShape]> { + let summary = "Reinterprets the element type of a memref."; + let description = [{ + Shape, layout and memory space are unchanged, so no data is moved: this is + purely a change of how the bits of each element are interpreted. + + Example: + + ```mlir + %0 = xtile.memref_bitcast %arg0 : memref<4x16xbf16> -> memref<4x16xi16> + ``` + }]; + + let arguments = (ins AnyMemRef:$source); + let results = (outs AnyMemRef:$result); + + let assemblyFormat = [{ + $source `:` type($source) `->` type($result) attr-dict + }]; + + let hasVerifier = 1; + let hasFolder = 1; +} + def ExtractTileOp : XTile_Op<"extract", [TiledBufferInterface]> { let summary = "Extract a tile from a memref."; let description = [{ diff --git a/third_party/xla/xla/ffi/BUILD b/third_party/xla/xla/ffi/BUILD index c383a0da055c20..e5137056eb2bca 100644 --- a/third_party/xla/xla/ffi/BUILD +++ b/third_party/xla/xla/ffi/BUILD @@ -31,14 +31,10 @@ package_group( ], ) -cc_library( +alias( name = "api", - hdrs = ["//xla/ffi/api:api_headers"], + actual = "//xla/ffi/api:api", visibility = [":ffi_internal"], - deps = [ - "//xla/ffi/api:c_api", - "@com_google_absl//absl/strings:string_view", - ], ) cc_library( diff --git a/third_party/xla/xla/ffi/api/BUILD b/third_party/xla/xla/ffi/api/BUILD index 287edfde68c27f..23c62b1f8aa99e 100644 --- a/third_party/xla/xla/ffi/api/BUILD +++ b/third_party/xla/xla/ffi/api/BUILD @@ -81,7 +81,7 @@ filegroup( cc_library( name = "api", hdrs = [":api_headers"], - visibility = ["//visibility:private"], + visibility = ["//xla/ffi:__pkg__"], deps = [ ":c_api", ], diff --git a/third_party/xla/xla/hlo/analysis/hlo_replication_analysis.cc b/third_party/xla/xla/hlo/analysis/hlo_replication_analysis.cc index 0f1e1400a6d4da..5443f988a4991e 100644 --- a/third_party/xla/xla/hlo/analysis/hlo_replication_analysis.cc +++ b/third_party/xla/xla/hlo/analysis/hlo_replication_analysis.cc @@ -499,11 +499,19 @@ absl::Status HloReplicationAnalysis::ComputeHloReplication() { HloReplication::UniqueOnAllDevices()); std::unique_ptr> sharding_tree = nullptr; - if (cross_partition_spmd_ && param->has_sharding()) { - ABSL_ASSIGN_OR_RETURN(auto result, - param->sharding().AsShapeTree(param->shape())); - sharding_tree = - std::make_unique>(std::move(result)); + if (cross_partition_spmd_) { + const HloSharding* sharding = nullptr; + if (module_->has_spmd_parameters_shardings()) { + CHECK_LT(i, module_->spmd_parameters_shardings().size()); + sharding = &module_->spmd_parameters_shardings()[i]; + } else if (param->has_sharding()) { + sharding = ¶m->sharding(); + } + if (sharding != nullptr) { + ABSL_ASSIGN_OR_RETURN(auto result, sharding->AsShapeTree(param->shape())); + sharding_tree = + std::make_unique>(std::move(result)); + } } const auto& replication = param->parameter_replicated_at_leaf_buffers(); @@ -513,7 +521,7 @@ absl::Status HloReplicationAnalysis::ComputeHloReplication() { if (!ShapeUtil::IsLeafIndex(param->shape(), index)) { return absl::OkStatus(); } - if (cross_partition_spmd_ && param->has_sharding()) { + if (sharding_tree != nullptr) { // In cross-partition spmd mode, set parameter replication status // based on the parameter's sharding. *shape_tree.mutable_element(index) = diff --git a/third_party/xla/xla/hlo/analysis/hlo_replication_analysis_test.cc b/third_party/xla/xla/hlo/analysis/hlo_replication_analysis_test.cc index 1b249702307944..6c601626141e72 100644 --- a/third_party/xla/xla/hlo/analysis/hlo_replication_analysis_test.cc +++ b/third_party/xla/xla/hlo/analysis/hlo_replication_analysis_test.cc @@ -19,6 +19,7 @@ limitations under the License. #include #include +#include #include #include "absl/types/span.h" #include "xla/hlo/ir/hlo_computation.h" @@ -223,6 +224,35 @@ ENTRY entry { FindInstruction(module.get(), "all-reduce-different-operand"), {})); } +TEST_F(HloReplicationAnalysisTest, + CrossPartitionSpmdWithModuleParameterShardings) { + const std::string module_str = R"( +HloModule CrossPartitionSpmdWithModuleParameterShardings + +ENTRY entry { + param = (f32[4096,4096]{1,0}, f32[4096,4096]{1,0}) + parameter(0), sharding={{maximal device=0}, {replicated}} + gte0 = f32[4096,4096]{1,0} get-tuple-element(param), index=0 + gte1 = f32[4096,4096]{1,0} get-tuple-element(param), index=1 + ROOT add = f32[4096,4096]{1,0} add(gte0, gte1) +} +)"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule( + module_str, /*replica_count=*/4)); + HloInstruction* param = module->entry_computation()->parameter_instruction(0); + module->set_spmd_parameters_shardings({param->sharding()}); + param->clear_sharding(); + + ASSERT_OK_AND_ASSIGN( + std::unique_ptr analysis, + HloReplicationAnalysis::Run(module.get(), /*cross_partition_spmd=*/true)); + EXPECT_FALSE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "gte0"), {})); + EXPECT_TRUE(analysis->HloInstructionIsReplicatedAt( + FindInstruction(module.get(), "gte1"), {})); +} + TEST_F(HloReplicationAnalysisTest, NestedCall) { const std::string module_str = R"( HloModule NestedCall diff --git a/third_party/xla/xla/hlo/analysis/symbolic_map.h b/third_party/xla/xla/hlo/analysis/symbolic_map.h index 95883e7b20f123..08e66cb75c4684 100644 --- a/third_party/xla/xla/hlo/analysis/symbolic_map.h +++ b/third_party/xla/xla/hlo/analysis/symbolic_map.h @@ -20,6 +20,7 @@ limitations under the License. #include #include #include +#include #include "absl/log/check.h" #include "absl/types/span.h" diff --git a/third_party/xla/xla/hlo/builder/BUILD b/third_party/xla/xla/hlo/builder/BUILD index 9087f57cfc2ca3..8eb58ff87a003e 100644 --- a/third_party/xla/xla/hlo/builder/BUILD +++ b/third_party/xla/xla/hlo/builder/BUILD @@ -209,6 +209,7 @@ xla_cc_test( "//xla:comparison_util", "//xla:debug_options_flags", "//xla:shape_util", + "//xla:shuffle", "//xla:tuple_tree", "//xla:util", "//xla:xla_data_proto_cc", diff --git a/third_party/xla/xla/hlo/builder/xla_builder.cc b/third_party/xla/xla/hlo/builder/xla_builder.cc index a672fa261726b7..7072aaf18ec7a0 100644 --- a/third_party/xla/xla/hlo/builder/xla_builder.cc +++ b/third_party/xla/xla/hlo/builder/xla_builder.cc @@ -3126,6 +3126,28 @@ absl::StatusOr XlaBuilder::RevInternal( return AddInstruction(std::move(instr), HloOpcode::kReverse, {operand}); } +XlaOp XlaBuilder::Shuffle(XlaOp operand, absl::Span dimensions, + const ShuffleMode& mode) { + return ReportErrorOrReturn([&]() -> absl::StatusOr { + ABSL_ASSIGN_OR_RETURN(const Shape* operand_shape, GetShapePtr(operand)); + ABSL_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferShuffleShape( + *operand_shape, dimensions, mode)); + return ShuffleInternal(shape, operand, dimensions, mode); + }); +} + +absl::StatusOr XlaBuilder::ShuffleInternal( + const Shape& shape, XlaOp operand, absl::Span dimensions, + const ShuffleMode& mode) { + HloInstructionProto instr; + *instr.mutable_shape() = shape.ToProto(); + for (int64_t dim : dimensions) { + instr.add_dimensions(dim); + } + *instr.mutable_shuffle_mode() = mode; + return AddInstruction(std::move(instr), HloOpcode::kShuffle, {operand}); +} + XlaOp XlaBuilder::Sort(absl::Span operands, XlaComputationId comparator, int64_t dimension, bool is_stable) { @@ -6672,6 +6694,11 @@ XlaOp Rev(const XlaOp operand, absl::Span dimensions) { return operand.builder()->Rev(operand, dimensions); } +XlaOp Shuffle(const XlaOp operand, absl::Span dimensions, + const ShuffleMode& mode) { + return operand.builder()->Shuffle(operand, dimensions, mode); +} + XlaOp Sort(absl::Span operands, const XlaComputation& comparator, int64_t dimension, bool is_stable) { return Sort(operands, operands[0].builder()->AddSubComputation(comparator), diff --git a/third_party/xla/xla/hlo/builder/xla_builder.h b/third_party/xla/xla/hlo/builder/xla_builder.h index 3c0e937a8c1088..d91e6592ebae4a 100644 --- a/third_party/xla/xla/hlo/builder/xla_builder.h +++ b/third_party/xla/xla/hlo/builder/xla_builder.h @@ -1111,6 +1111,12 @@ class XlaBuilder { virtual absl::StatusOr RevInternal( const Shape& shape, XlaOp operand, absl::Span dimensions); + XlaOp Shuffle(XlaOp operand, absl::Span dimensions, + const ShuffleMode& mode); + virtual absl::StatusOr ShuffleInternal( + const Shape& shape, XlaOp operand, absl::Span dimensions, + const ShuffleMode& mode); + XlaOp Sort(absl::Span operands, XlaComputationId comparator, int64_t dimension = -1, bool is_stable = false); virtual absl::StatusOr SortInternal(const Shape& shape, @@ -1991,6 +1997,8 @@ class XlaBuilder { friend XlaOp Neg(XlaOp operand); friend XlaOp Transpose(XlaOp operand, absl::Span permutation); friend XlaOp Rev(XlaOp operand, absl::Span dimensions); + friend XlaOp Shuffle(XlaOp operand, absl::Span dimensions, + const ShuffleMode& mode); friend XlaOp Sort(absl::Span operands, const XlaComputation& comparator, int64_t dimension, bool is_stable); @@ -3413,6 +3421,15 @@ XlaOp Transpose(XlaOp operand, absl::Span permutation); // is moved to index dimension_size - 1 - i). XlaOp Rev(XlaOp operand, absl::Span dimensions); +// Enqueues a shuffle instruction onto the computation. The elements are +// shuffled along the given dimensions following the pattern selected by `mode`, +// which also carries the attributes of that mode: +// - rotate: the elements are (left) rotated by `shifts` along the given +// dimensions (i.e. the element at index i in the output is taken from index +// (i + shifts[j]) % dimension_size[j] of the operand). +XlaOp Shuffle(XlaOp operand, absl::Span dimensions, + const ShuffleMode& mode); + // Enqueues a sort instruction onto the computation, using 'comparator' for // comparisons. 'comparator' needs to define a strict weak order. 'is_stable' // determines whether the stable sorting should be used. diff --git a/third_party/xla/xla/hlo/builder/xla_builder_test.cc b/third_party/xla/xla/hlo/builder/xla_builder_test.cc index 4aa44b97be8114..05b30faa5a41e8 100644 --- a/third_party/xla/xla/hlo/builder/xla_builder_test.cc +++ b/third_party/xla/xla/hlo/builder/xla_builder_test.cc @@ -59,6 +59,7 @@ limitations under the License. #include "xla/service/pattern_matcher.h" #include "xla/shape.h" #include "xla/shape_util.h" +#include "xla/shuffle.h" #include "xla/tsl/platform/statusor.h" #include "xla/tuple_tree.h" #include "xla/util.h" @@ -3586,6 +3587,20 @@ TEST(XlaBuilderTest, UnboundedReverse) { GmockMatch(m::Op().WithShapeEqualTo(&expected))); } +TEST(XlaBuilderTest, UnboundedShuffle) { + XlaBuilder b(TestName()); + ASSERT_OK_AND_ASSIGN(const Shape operand, ParseShape("f32[?, 10]")); + ASSERT_OK_AND_ASSIGN(const Shape expected, ParseShape("f32[?, 10]")); + + Shuffle(Parameter(&b, 0, operand, "operand"), /*dimensions=*/{0, 1}, + shuffle::Rotate(/*shifts=*/{1, 3})); + ASSERT_OK_AND_ASSIGN(const std::unique_ptr module, + BuildHloModule(b)); + + EXPECT_THAT(GetRoot(*module), + GmockMatch(m::Op().WithShapeEqualTo(&expected))); +} + TEST(XlaBuilderTest, UnboundedRngBitGenerator) { XlaBuilder b(TestName()); TF_ASSERT_OK_AND_ASSIGN(const Shape initial_state, ParseShape("u32[?, 10]")); diff --git a/third_party/xla/xla/hlo/ir/BUILD b/third_party/xla/xla/hlo/ir/BUILD index 9e44d36b13c01c..a25fbd8777f887 100644 --- a/third_party/xla/xla/hlo/ir/BUILD +++ b/third_party/xla/xla/hlo/ir/BUILD @@ -63,7 +63,7 @@ cc_library( ], visibility = internal_visibility([ ":friends", - "//xla/backends/gpu/libraries/native_custom_call_thunks:handler_allowlist", + "//xla:native_custom_call_handler_allowlist", ]), deps = [ ":backend_config", diff --git a/third_party/xla/xla/hlo/ir/dfs_hlo_visitor.h b/third_party/xla/xla/hlo/ir/dfs_hlo_visitor.h index 592af8fb1d8553..efd83fd73ba18e 100644 --- a/third_party/xla/xla/hlo/ir/dfs_hlo_visitor.h +++ b/third_party/xla/xla/hlo/ir/dfs_hlo_visitor.h @@ -318,6 +318,7 @@ class DfsHloVisitorBase { virtual absl::Status HandleScan(HloInstructionPtr hlo) = 0; virtual absl::Status HandleScatter(HloInstructionPtr hlo) = 0; virtual absl::Status HandleSelectAndScatter(HloInstructionPtr hlo) = 0; + virtual absl::Status HandleShuffle(HloInstructionPtr hlo) = 0; virtual absl::Status HandleSlice(HloInstructionPtr hlo) = 0; virtual absl::Status HandleSort(HloInstructionPtr hlo) = 0; virtual absl::Status HandleTranspose(HloInstructionPtr hlo) = 0; diff --git a/third_party/xla/xla/hlo/ir/dfs_hlo_visitor_with_default.h b/third_party/xla/xla/hlo/ir/dfs_hlo_visitor_with_default.h index 24177f47f7be54..b691eb7ed027a0 100644 --- a/third_party/xla/xla/hlo/ir/dfs_hlo_visitor_with_default.h +++ b/third_party/xla/xla/hlo/ir/dfs_hlo_visitor_with_default.h @@ -170,6 +170,9 @@ class DfsHloVisitorWithDefaultBase absl::Status HandleReverse(HloInstructionPtr reverse) override { return DefaultAction(reverse); } + absl::Status HandleShuffle(HloInstructionPtr shuffle) override { + return DefaultAction(shuffle); + } absl::Status HandleSort(HloInstructionPtr sort) override { return DefaultAction(sort); } diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction.cc b/third_party/xla/xla/hlo/ir/hlo_instruction.cc index d21bebe569b702..273ace08d474d4 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instruction.cc +++ b/third_party/xla/xla/hlo/ir/hlo_instruction.cc @@ -588,6 +588,13 @@ absl::StatusOr> HloInstruction::CreateFromProto( std::vector(proto.dimensions().begin(), proto.dimensions().end())); break; + case HloOpcode::kShuffle: + instruction = + CreateShuffle(shape, operands(0), + std::vector(proto.dimensions().begin(), + proto.dimensions().end()), + proto.shuffle_mode()); + break; case HloOpcode::kConcatenate: TF_RET_CHECK(proto.dimensions().size() == 1) << "Concatenate instruction should have 1 dimension but sees " @@ -2157,6 +2164,13 @@ HloInstruction::CreateCollectivePermuteStart( return std::make_unique(shape, operand, dimensions); } +/* static */ std::unique_ptr HloInstruction::CreateShuffle( + const Shape& shape, HloInstruction* operand, + absl::Span dimensions, const ShuffleMode& mode) { + return std::make_unique(shape, operand, dimensions, + mode); +} + /* static */ std::unique_ptr HloInstruction::CreateAfterAll( absl::Span operands) { CHECK(!operands.empty()); @@ -2855,6 +2869,7 @@ std::unique_ptr HloInstruction::CloneWithNewOperands( case HloOpcode::kRecv: case HloOpcode::kRecvDone: case HloOpcode::kReverse: + case HloOpcode::kShuffle: case HloOpcode::kConcatenate: case HloOpcode::kReduce: case HloOpcode::kTranspose: @@ -3570,6 +3585,7 @@ bool HloInstruction::IdenticalSlowPath( case HloOpcode::kTriangularSolve: case HloOpcode::kCholesky: case HloOpcode::kTopK: + case HloOpcode::kShuffle: LOG(FATAL) << "Base class impl called for opcode with subclass: " << opcode(); } @@ -5102,6 +5118,8 @@ absl::Status HloInstruction::Visit( return visitor->HandleTranspose(this); case HloOpcode::kReverse: return visitor->HandleReverse(this); + case HloOpcode::kShuffle: + return visitor->HandleShuffle(this); case HloOpcode::kReducePrecision: return visitor->HandleReducePrecision(this); case HloOpcode::kSlice: @@ -5421,6 +5439,7 @@ static UseKind OperandElementUse(const HloInstruction& instr, case HloOpcode::kSlice: case HloOpcode::kTranspose: case HloOpcode::kGather: + case HloOpcode::kShuffle: return UseKind::kUse; case HloOpcode::kPad: // Pad reuses the padding value but not the padded array elements. @@ -5608,6 +5627,15 @@ std::string PrecisionToString(const PrecisionConfig::Precision& precision) { return absl::AsciiStrToLower(PrecisionConfig::Precision_Name(precision)); } +std::string ShuffleModeToString(ShuffleMode::ModeCase shuffle_mode) { + switch (shuffle_mode) { + case ShuffleMode::kRotate: + return "rotate"; + case ShuffleMode::MODE_NOT_SET: + return "invalid"; + } +} + template void AbslStringify(Sink& sink, const ResultAccuracy::Tolerance& tolerance) { absl::Format(&sink, "tolerance={atol=%v,rtol=%v,ulps=%v}", tolerance.atol(), @@ -5870,6 +5898,16 @@ absl::StatusOr StringToPrecision( "precision"); } +absl::StatusOr StringToShuffleMode( + absl::string_view mode) { + if (mode == "rotate") { + return ShuffleMode::ModeCase::kRotate; + } + return InvalidArgument("Unknown shuffle mode: %s", mode); + // return StringToEnum(mode, ShuffleModeToString, + // "shuffle mode"); +} + absl::StatusOr StringToResultAccuracy( absl::string_view name) { return StringToEnum(name, ResultAccuracyToString, diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction.h b/third_party/xla/xla/hlo/ir/hlo_instruction.h index aa10a9def8c674..a899406d8b6142 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instruction.h +++ b/third_party/xla/xla/hlo/ir/hlo_instruction.h @@ -1214,6 +1214,12 @@ class HloInstruction { const Shape& shape, HloInstruction* operand, absl::Span dimensions); + // Creates a shuffle instruction, which shuffles the elements of `operand` + // along the given dimensions following the pattern selected by `mode`. + static std::unique_ptr CreateShuffle( + const Shape& shape, HloInstruction* operand, + absl::Span dimensions, const ShuffleMode& mode); + // Creates a Afterall instruction used for joining or creating new values of // token type which thread through side-effecting operations. Operands must // all be tokens, calls without operands generates a token. @@ -2955,6 +2961,7 @@ std::string ResultAccuracyToleranceToString( std::string RandomAlgorithmToString(const RandomAlgorithm& algorithm); std::string RandomDistributionToString(const RandomDistribution& distribution); std::string PrecisionToString(const PrecisionConfig::Precision& precision); +std::string ShuffleModeToString(ShuffleMode::ModeCase shuffle_mode); std::string ResultAccuracyToString(ResultAccuracy::Mode accuracy_mode); std::string AlgorithmToString(const PrecisionConfig::Algorithm& algorithm); std::string DotDimensionNumbersToString(const DotDimensionNumbers& dnums); @@ -2973,6 +2980,8 @@ absl::StatusOr StringToRandomDistribution( const std::string& name); absl::StatusOr StringToPrecision( const std::string& name); +absl::StatusOr StringToShuffleMode( + absl::string_view mode); absl::StatusOr StringToAlgorithm( const std::string& name); absl::StatusOr StringToResultAccuracy( diff --git a/third_party/xla/xla/hlo/ir/hlo_instructions.cc b/third_party/xla/xla/hlo/ir/hlo_instructions.cc index ce9d59cefa5a12..4071f0c416a838 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instructions.cc +++ b/third_party/xla/xla/hlo/ir/hlo_instructions.cc @@ -1984,6 +1984,59 @@ std::unique_ptr HloReduceInstruction::CloneWithNewOperandsImpl( dimensions(), to_apply()); } +HloShuffleInstruction::HloShuffleInstruction( + const Shape& shape, HloInstruction* operand, + absl::Span dimensions, const ShuffleMode& mode) + : HloDimensionsInstruction(HloOpcode::kShuffle, shape, dimensions), + mode_(mode) { + AppendOperand(operand); +} + +bool HloShuffleInstruction::IdenticalSlowPath( + const HloInstruction& other, + absl::FunctionRef + eq_computations) const { + const auto& casted_other = static_cast(other); + // Shuffle results are determined by the shuffled dimensions and by the mode + // together with its attributes. + return dimensions() == casted_other.dimensions() && + protobuf_util::HaveSameSerialization(mode_, casted_other.mode_); +} + +std::unique_ptr HloShuffleInstruction::CloneWithNewOperandsImpl( + const Shape& shape, absl::Span new_operands, + HloCloneContext* context) const { + CHECK_EQ(new_operands.size(), 1); + return std::make_unique(shape, new_operands[0], + dimensions(), mode_); +} + +void HloShuffleInstruction::PrintExtraAttributesImpl( + AttributePrinter& printer, const HloPrintOptions& options) const { + HloDimensionsInstruction::PrintExtraAttributesImpl(printer, options); + printer.Next([this](Printer* printer) { + printer->Append("mode="); + printer->Append(ShuffleModeToString(mode())); + }); + // Each mode's inner attributes. + switch (mode()) { + case ShuffleMode::kRotate: + printer.Next([this](Printer* printer) { + printer->Append("shifts={"); + AppendJoin(printer, rotate().shifts(), ","); + printer->Append("}"); + }); + break; + case ShuffleMode::MODE_NOT_SET: + break; + } +} + +void HloShuffleInstruction::ToProto(HloInstructionProto* proto) const { + HloDimensionsInstruction::ToProto(proto); + *proto->mutable_shuffle_mode() = mode_; +} + HloScanInstruction::HloScanInstruction(const Shape& shape, absl::Span inputs, absl::Span inits, diff --git a/third_party/xla/xla/hlo/ir/hlo_instructions.h b/third_party/xla/xla/hlo/ir/hlo_instructions.h index 054e036d666b7d..f67a1714749161 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instructions.h +++ b/third_party/xla/xla/hlo/ir/hlo_instructions.h @@ -72,6 +72,7 @@ class HloDimensionsInstruction : public HloInstruction { case HloOpcode::kConcatenate: case HloOpcode::kReduce: case HloOpcode::kReverse: + case HloOpcode::kShuffle: case HloOpcode::kSort: case HloOpcode::kTranspose: return true; @@ -1178,6 +1179,47 @@ class HloReverseInstruction : public HloDimensionsInstruction { HloCloneContext* context) const override; }; +class HloShuffleInstruction : public HloDimensionsInstruction { + public: + explicit HloShuffleInstruction(const Shape& shape, HloInstruction* operand, + absl::Span dimensions, + const ShuffleMode& mode); + + static bool ClassOf(const HloInstruction* hlo) { + return hlo->opcode() == HloOpcode::kShuffle; + } + + const ShuffleMode& shuffle_mode() const { return mode_; } + ShuffleMode* mutable_shuffle_mode() { return &mode_; } + ShuffleMode::ModeCase mode() const { return mode_.mode_case(); } + void ToProto(HloInstructionProto* proto) const override; + + // Accessors for the inner attributes of each mode. + const ShuffleMode::Rotate& rotate() const { + CHECK(mode_.has_rotate()); + return mode_.rotate(); + } + ShuffleMode::Rotate* mutable_rotate() { + CHECK(mode_.has_rotate()); + return mode_.mutable_rotate(); + } + + private: + std::unique_ptr CloneWithNewOperandsImpl( + const Shape& shape, absl::Span new_operands, + HloCloneContext* context) const override; + + bool IdenticalSlowPath( + const HloInstruction& other, + absl::FunctionRef + eq_computations) const override; + + void PrintExtraAttributesImpl(AttributePrinter& printer, + const HloPrintOptions& options) const override; + + ShuffleMode mode_; +}; + class HloConcatenateInstruction : public HloDimensionsInstruction { public: explicit HloConcatenateInstruction(const Shape& shape, diff --git a/third_party/xla/xla/hlo/ir/hlo_opcode.h b/third_party/xla/xla/hlo/ir/hlo_opcode.h index b25434781f4158..9da5dc0342a2d7 100644 --- a/third_party/xla/xla/hlo/ir/hlo_opcode.h +++ b/third_party/xla/xla/hlo/ir/hlo_opcode.h @@ -165,6 +165,7 @@ namespace xla { V(kShiftLeft, "shift-left", 2) \ V(kShiftRightArithmetic, "shift-right-arithmetic", 2) \ V(kShiftRightLogical, "shift-right-logical", 2) \ + V(kShuffle, "shuffle", 1) \ V(kSign, "sign", 1) \ V(kSin, "sine", 1) \ V(kSinh, "sinh", 1) \ diff --git a/third_party/xla/xla/hlo/parser/BUILD b/third_party/xla/xla/hlo/parser/BUILD index 6b90eb2a45b9d3..1d0bbd1a6dd6b4 100644 --- a/third_party/xla/xla/hlo/parser/BUILD +++ b/third_party/xla/xla/hlo/parser/BUILD @@ -33,6 +33,7 @@ cc_library( "//xla:literal_util", "//xla:shape_layout", "//xla:shape_util", + "//xla:shuffle", "//xla:tuple_tree", "//xla:types", "//xla:util", diff --git a/third_party/xla/xla/hlo/parser/hlo_parser.cc b/third_party/xla/xla/hlo/parser/hlo_parser.cc index 518f3aee55f2c2..268eedd5fbd1e9 100644 --- a/third_party/xla/xla/hlo/parser/hlo_parser.cc +++ b/third_party/xla/xla/hlo/parser/hlo_parser.cc @@ -87,6 +87,7 @@ limitations under the License. #include "xla/shape.h" #include "xla/shape_layout.h" #include "xla/shape_util.h" +#include "xla/shuffle.h" #include "xla/tsl/lib/gtl/map_util.h" #include "xla/tsl/platform/logging.h" #include "xla/tuple_tree.h" @@ -200,6 +201,7 @@ bool CanInferShape(HloOpcode code) { case HloOpcode::kShiftLeft: case HloOpcode::kShiftRightArithmetic: case HloOpcode::kShiftRightLogical: + case HloOpcode::kShuffle: case HloOpcode::kSign: case HloOpcode::kSin: case HloOpcode::kSinh: @@ -344,6 +346,7 @@ class HloParserImpl : public HloParser { kShapeList, kEnum, kRandomAlgorithm, + kShuffleMode, kPrecisionAlgorithm, kResultAccuracyType, kAliasing, @@ -636,6 +639,7 @@ class HloParserImpl : public HloParser { bool ParseRandomDistribution(RandomDistribution* result); bool ParseConvKind(ConvolutionKind* result); bool ParseRandomAlgorithm(RandomAlgorithm* result); + bool ParseShuffleMode(ShuffleMode::ModeCase* result); bool ParsePrecision(PrecisionConfig::Precision* result); bool ParseAlgorithm(PrecisionConfig::Algorithm* result); bool ParseResultAccuracyType(ResultAccuracy::Mode* result); @@ -3137,6 +3141,43 @@ HloInstruction* HloParserImpl::CreateInstruction( // NOLINT return builder->AddInstruction( HloInstruction::CreateReverse(*shape, operands[0], *dimensions)); } + case HloOpcode::kShuffle: { + optional> dimensions; + attrs["dimensions"] = {/*required=*/true, AttrTy::kBracedInt64List, + &dimensions}; + optional mode; + attrs["mode"] = {/*required=*/true, AttrTy::kShuffleMode, &mode}; + // Attributes below are specific to a subset of the modes, so they are + // optional here and checked against the mode afterwards. + optional> shifts; + attrs["shifts"] = {/*required=*/false, AttrTy::kBracedInt64List, &shifts}; + if ((!preset_operands && + !ParseOperands(&operands, builder, /*expected_size=*/1)) || + !ParseAttributes(attrs, allow_attributes, shape)) { + return nullptr; + } + ShuffleMode shuffle_mode; + switch (*mode) { + case ShuffleMode::kRotate: + if (!shifts.has_value()) { + TokenError("expects shifts for a shuffle in rotate mode"); + return nullptr; + } + shuffle_mode = shuffle::Rotate(*shifts); + break; + case ShuffleMode::MODE_NOT_SET: + TokenError("expects a shuffle mode"); + return nullptr; + } + if (!maybe_infer_shape([&] { + return ShapeInference::InferShuffleShape(operands[0]->shape(), + *dimensions, shuffle_mode); + })) { + return nullptr; + } + return builder->AddInstruction(HloInstruction::CreateShuffle( + *shape, operands[0], *dimensions, shuffle_mode)); + } case HloOpcode::kSelectAndScatter: { optional select; attrs["select"] = {/*required=*/true, AttrTy::kHloComputation, &select}; @@ -6416,6 +6457,15 @@ bool HloParserImpl::ParseAttributeHelper( static_cast*>(attr_out_ptr)->emplace(result); return true; } + case AttrTy::kShuffleMode: { + ShuffleMode::ModeCase result; + if (!ParseShuffleMode(&result)) { + return false; + } + static_cast*>(attr_out_ptr) + ->emplace(result); + return true; + } case AttrTy::kPrecisionAlgorithm: { PrecisionConfig::Algorithm result; if (!ParseAlgorithm(&result)) { @@ -8560,6 +8610,22 @@ bool HloParserImpl::ParseRandomAlgorithm(RandomAlgorithm* result) { return true; } +bool HloParserImpl::ParseShuffleMode(ShuffleMode::ModeCase* result) { + VLOG(kDebugLevel) << "ParseShuffleMode"; + if (lexer_.GetKind() != TokKind::kIdent) { + return TokenError("expects shuffle mode"); + } + std::string val = lexer_.GetStrVal(); + auto status_or_result = StringToShuffleMode(val); + if (!status_or_result.ok()) { + return TokenError(StrFormat("expects shuffle mode but sees: %s, error: %s", + val, status_or_result.status().message())); + } + *result = status_or_result.value(); + lexer_.Lex(); + return true; +} + bool HloParserImpl::ParsePrecision(PrecisionConfig::Precision* result) { VLOG(kDebugLevel) << "ParsePrecision"; if (lexer_.GetKind() != TokKind::kIdent) { diff --git a/third_party/xla/xla/hlo/parser/hlo_parser_test.cc b/third_party/xla/xla/hlo/parser/hlo_parser_test.cc index 9ff45bd6026930..6502d1f3d52d87 100644 --- a/third_party/xla/xla/hlo/parser/hlo_parser_test.cc +++ b/third_party/xla/xla/hlo/parser/hlo_parser_test.cc @@ -775,6 +775,18 @@ ENTRY %Reverse4DFloatArrayOnDim01.v2 () -> f32[4,3,2,1] { ROOT %reverse = f32[4,3,2,1]{0,1,2,3} reverse(f32[4,3,2,1]{0,1,2,3} %constant), dimensions={0,1} } +)" +}, +// shuffle(constant, mode=rotate) +{ +"ShuffleRotate2D", +R"(HloModule ShuffleRotate2DFloatArrayOnDim01_module, entry_computation_layout={()->f32[4,3]{1,0}} + +ENTRY %ShuffleRotate2DFloatArrayOnDim01.v2 () -> f32[4,3] { + %constant = f32[4,3]{1,0} constant({ { 1, 2, 3 }, { 4, 5, 6 }, { 7, 8, 9 }, { 10, 11, 12 } }) + ROOT %shuffle = f32[4,3]{1,0} shuffle(f32[4,3]{1,0} %constant), dimensions={0,1}, mode=rotate, shifts={1,2} +} + )" }, // concat @@ -8720,5 +8732,194 @@ ENTRY %main (p0: f32[]) -> f32[] { "raw_computation_config"); EXPECT_EQ(roundtrip_module->ToString(), printed); } + +TEST_F(HloParserTest, ShuffleUnsupportedMode) { + const std::string original = R"(HloModule shuffle_unsupported_mode + +ENTRY %entry (p: f32[4,3]) -> f32[4,3] { + %p = f32[4,3]{1,0} parameter(0) + ROOT %shuffle = f32[4,3]{1,0} shuffle(f32[4,3]{1,0} %p), dimensions={0,1}, mode=transpose +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr(result.status().message(), + "expects shuffle mode but sees: transpose, error: Unknown " + "shuffle mode: transpose"); +} + +TEST_F(HloParserTest, ShuffleModeIsCaseSensitive) { + const std::string original = R"(HloModule shuffle_uppercase_mode + +ENTRY %entry (p: f32[4,3]) -> f32[4,3] { + %p = f32[4,3]{1,0} parameter(0) + ROOT %shuffle = f32[4,3]{1,0} shuffle(f32[4,3]{1,0} %p), dimensions={0,1}, mode=ROTATE, shifts={1,2} +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr(result.status().message(), + "expects shuffle mode but sees: ROTATE, error: Unknown " + "shuffle mode: ROTATE"); +} + +TEST_F(HloParserTest, ShuffleModeIsNotAnIdentifier) { + const std::string original = R"(HloModule shuffle_non_identifier_mode + +ENTRY %entry (p: f32[4,3]) -> f32[4,3] { + %p = f32[4,3]{1,0} parameter(0) + ROOT %shuffle = f32[4,3]{1,0} shuffle(f32[4,3]{1,0} %p), dimensions={0,1}, mode=0, shifts={1,2} +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr(result.status().message(), "expects shuffle mode"); +} + +TEST_F(HloParserTest, ShuffleMissingMode) { + const std::string original = R"(HloModule shuffle_missing_mode + +ENTRY %entry (p: f32[4,3]) -> f32[4,3] { + %p = f32[4,3]{1,0} parameter(0) + ROOT %shuffle = f32[4,3]{1,0} shuffle(f32[4,3]{1,0} %p), dimensions={0,1}, shifts={1,2} +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr(result.status().message(), + "attribute mode is expected but not seen"); +} + +TEST_F(HloParserTest, ShuffleMissingDimensions) { + const std::string original = R"(HloModule shuffle_missing_dimensions + +ENTRY %entry (p: f32[4,3]) -> f32[4,3] { + %p = f32[4,3]{1,0} parameter(0) + ROOT %shuffle = f32[4,3]{1,0} shuffle(f32[4,3]{1,0} %p), mode=rotate, shifts={1,2} +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr(result.status().message(), + "attribute dimensions is expected but not seen"); +} + +// The `shifts` attribute is specific to the rotate mode, so it must be present +// if and only if the mode is rotate. +TEST_F(HloParserTest, ShuffleRotateWithoutShifts) { + const std::string original = R"(HloModule shuffle_rotate_without_shifts + +ENTRY %entry (p: f32[4,3]) -> f32[4,3] { + %p = f32[4,3]{1,0} parameter(0) + ROOT %shuffle = f32[4,3]{1,0} shuffle(f32[4,3]{1,0} %p), dimensions={0,1}, mode=rotate +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr(result.status().message(), + "expects shifts for a shuffle in rotate mode"); +} + +TEST_F(HloParserTest, ShuffleUnknownAttribute) { + const std::string original = R"(HloModule shuffle_unknown_attribute + +ENTRY %entry (p: f32[4,3]) -> f32[4,3] { + %p = f32[4,3]{1,0} parameter(0) + ROOT %shuffle = f32[4,3]{1,0} shuffle(f32[4,3]{1,0} %p), dimensions={0,1}, mode=rotate, shifts={1,2}, strides={1,1} +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr(result.status().message(), + "unexpected attribute \"strides\""); +} + +TEST_F(HloParserTest, ShuffleRepeatedAttribute) { + const std::string original = R"(HloModule shuffle_repeated_attribute + +ENTRY %entry (p: f32[4,3]) -> f32[4,3] { + %p = f32[4,3]{1,0} parameter(0) + ROOT %shuffle = f32[4,3]{1,0} shuffle(f32[4,3]{1,0} %p), dimensions={0,1}, mode=rotate, shifts={1,2}, shifts={1,2} +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr(result.status().message(), "attribute shifts already exists"); +} + +TEST_F(HloParserTest, ShuffleWrongOperandsSize) { + const std::string original = R"(HloModule shuffle_wrong_operands_size + +ENTRY %entry (p: f32[4,3]) -> f32[4,3] { + %p = f32[4,3]{1,0} parameter(0) + ROOT %shuffle = f32[4,3]{1,0} shuffle(f32[4,3]{1,0} %p, f32[4,3]{1,0} %p), dimensions={0,1}, mode=rotate, shifts={1,2} +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr(result.status().message(), "expects 1 operands"); +} + +// The tests below omit the shape of the shuffle, which makes the parser infer +// it and report the shape inference errors that the mode violates. +TEST_F(HloParserTest, ShuffleShiftsSizeMismatchesDimensionsSize) { + const std::string original = R"(HloModule shuffle_shifts_size_mismatch + +ENTRY %entry (p: f32[4,3]) -> f32[4,3] { + %p = f32[4,3]{1,0} parameter(0) + ROOT %shuffle = shuffle(f32[4,3]{1,0} %p), dimensions={0,1}, mode=rotate, shifts={1} +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr(result.status().message(), + "failed to infer shape for opcode: shuffle, error: " + "dimensions and shifts must have the same size, got 2 and 1"); +} + +TEST_F(HloParserTest, ShuffleNoDimensions) { + const std::string original = R"(HloModule shuffle_no_dimensions + +ENTRY %entry (p: f32[4,3]) -> f32[4,3] { + %p = f32[4,3]{1,0} parameter(0) + ROOT %shuffle = shuffle(f32[4,3]{1,0} %p), dimensions={}, mode=rotate, shifts={} +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr(result.status().message(), + "A shuffle must shuffle at least one dimension."); +} + +TEST_F(HloParserTest, ShuffleDuplicateDimensions) { + const std::string original = R"(HloModule shuffle_duplicate_dimensions + +ENTRY %entry (p: f32[4,3]) -> f32[4,3] { + %p = f32[4,3]{1,0} parameter(0) + ROOT %shuffle = shuffle(f32[4,3]{1,0} %p), dimensions={0,0}, mode=rotate, shifts={1,2} +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr(result.status().message(), + "A dimension number is duplicated in shuffle."); +} + +TEST_F(HloParserTest, ShuffleDimensionOutOfBounds) { + const std::string original = R"(HloModule shuffle_dimension_out_of_bounds + +ENTRY %entry (p: f32[4,3]) -> f32[4,3] { + %p = f32[4,3]{1,0} parameter(0) + ROOT %shuffle = shuffle(f32[4,3]{1,0} %p), dimensions={0,2}, mode=rotate, shifts={1,2} +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr( + result.status().message(), + "One of the shuffle dimensions (2) is out-of-bounds in shape f32[4,3]"); +} + +TEST_F(HloParserTest, ShuffleOperandIsNotAnArray) { + const std::string original = R"(HloModule shuffle_non_array_operand + +ENTRY %entry (p: f32[4,3]) -> (f32[4,3]) { + %p = f32[4,3]{1,0} parameter(0) + %tuple = (f32[4,3]{1,0}) tuple(f32[4,3]{1,0} %p) + ROOT %shuffle = shuffle((f32[4,3]{1,0}) %tuple), dimensions={0}, mode=rotate, shifts={1} +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr(result.status().message(), + "Expected array argument for operand of shuffle"); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/hlo/tools/hlo_diff/hlo_diff_summary_test.cc b/third_party/xla/xla/hlo/tools/hlo_diff/hlo_diff_summary_test.cc index 9002ae7455f966..93aa4055b6f1bf 100644 --- a/third_party/xla/xla/hlo/tools/hlo_diff/hlo_diff_summary_test.cc +++ b/third_party/xla/xla/hlo/tools/hlo_diff/hlo_diff_summary_test.cc @@ -506,7 +506,7 @@ subtract.0 = f32[] subtract(constant.0, constant.1) /*left_unmatched_instruction_count=*/3, /*right_unmatched_instruction_count=*/0)), FieldsAre( - /*fingerprint=*/2080109884842586091U, + /*fingerprint=*/15551884552649592524U, /*computation_groups=*/ UnorderedElementsAre(FieldsAre( /*left_computations=*/IsEmpty(), diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/BUILD b/third_party/xla/xla/hlo/transforms/simplifiers/BUILD index 1f3f0aa3423a74..7fd054c9e50849 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/BUILD +++ b/third_party/xla/xla/hlo/transforms/simplifiers/BUILD @@ -486,6 +486,7 @@ cc_library( "//xla:literal_util", "//xla:permutation_util", "//xla:shape_util", + "//xla:shuffle", "//xla:status_macros", "//xla:util", "//xla:window_util", @@ -508,6 +509,7 @@ cc_library( "//xla/service:shape_inference", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base:nullability", + "@com_google_absl//absl/container:btree", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/container:inlined_vector", @@ -1217,7 +1219,6 @@ cc_library( "@com_google_absl//absl/strings:str_format", "@com_google_absl//absl/strings:string_view", "@com_google_absl//absl/time", - "@tsl//tsl/platform:errors", ], ) @@ -1225,6 +1226,7 @@ xla_cc_test( name = "hlo_constant_folding_test", srcs = ["hlo_constant_folding_test.cc"], deps = [ + ":flatten_call_graph", ":hlo_constant_folding", "//xla:literal", "//xla:literal_util", @@ -1233,12 +1235,14 @@ xla_cc_test( "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", "//xla/hlo/parser:hlo_parser", + "//xla/hlo/testlib:filecheck", "//xla/hlo/testlib:hlo_hardware_independent_test_base", "//xla/hlo/testlib:pattern_matcher_gmock", "//xla/hlo/testlib:test", "//xla/hlo/utils:hlo_matchers", "//xla/service:pattern_matcher", "//xla/tsl/platform:statusor", + "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/strings:string_view", "@com_google_absl//absl/types:span", "@com_google_googletest//:gtest_main", diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc index 97edf369ef7ba6..303b102cd71dde 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc @@ -32,6 +32,7 @@ limitations under the License. #include "absl/algorithm/container.h" #include "absl/base/nullability.h" +#include "absl/container/btree_map.h" #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" #include "absl/container/inlined_vector.h" @@ -77,6 +78,7 @@ limitations under the License. #include "xla/service/shape_inference.h" #include "xla/shape.h" #include "xla/shape_util.h" +#include "xla/shuffle.h" #include "xla/status_macros.h" #include "xla/util.h" #include "xla/window_util.h" @@ -7098,6 +7100,84 @@ absl::Status AlgebraicSimplifierVisitor::HandleReverse( return absl::OkStatus(); } +absl::Status AlgebraicSimplifierVisitor::HandleShuffle(HloInstruction* hlo) { + auto* shuffle = Cast(hlo); + + switch (shuffle->mode()) { + case ShuffleMode::kRotate: { + // Accumulate shifts per dimension (in ascending order of dim). + absl::btree_map combined_shifts; + for (int64_t i = 0; i < shuffle->dimensions().size(); ++i) { + combined_shifts[shuffle->dimensions()[i]] += + shuffle->rotate().shifts(i); + } + + // rotate(rotate(x, {d}, {m}), {d}, {n}) ==> rotate(x, {d}, {m + n}) + HloInstruction* base_operand = shuffle->mutable_operand(0); + while (base_operand->opcode() == HloOpcode::kShuffle && + Cast(base_operand)->mode() == + ShuffleMode::kRotate) { + auto* inner_shuffle = Cast(base_operand); + for (int64_t i = 0; i < inner_shuffle->dimensions().size(); ++i) { + combined_shifts[inner_shuffle->dimensions()[i]] += + inner_shuffle->rotate().shifts(i); + } + base_operand = inner_shuffle->mutable_operand(0); + } + + // Helper to check if base_operand is a splat along dim. + auto is_dim_splat = [&](int64_t dim) -> bool { + if (base_operand->IsConstant() && + base_operand->literal().IsAllFirst()) { + return true; + } + if (base_operand->opcode() == HloOpcode::kBroadcast) { + return !absl::c_linear_search(base_operand->dimensions(), dim); + } + return false; + }; + + // Canonicalize shifts & dimensions, remove no-op dims. + DimensionVector new_dimensions; + DimensionVector new_shifts; + for (const auto& [dim, total_shift] : combined_shifts) { + int64_t dim_size = shuffle->shape().dimensions(dim); + // No-op: size 1 dim + if (dim_size <= 1) { + continue; + } + // No-op: splat along dim + if (is_dim_splat(dim)) { + continue; + } + int64_t norm_shift = shuffle::NormalizeShift(total_shift, dim_size); + // No-op: shift == 0 + if (norm_shift == 0) { + continue; + } + new_dimensions.push_back(dim); + new_shifts.push_back(norm_shift); + } + + // Replace if changed. + if (new_dimensions.empty()) { + return ReplaceInstruction(shuffle, base_operand); + } + if (base_operand != shuffle->operand(0) || + new_dimensions != shuffle->dimensions() || + !absl::c_equal(new_shifts, shuffle->rotate().shifts())) { + auto new_shuffle = HloInstruction::CreateShuffle( + shuffle->shape(), base_operand, new_dimensions, + shuffle::Rotate(new_shifts)); + return ReplaceWithNewInstruction(shuffle, std::move(new_shuffle)); + } + return absl::OkStatus(); + } + default: + return absl::OkStatus(); + } +} + absl::StatusOr AlgebraicSimplifierVisitor::TrySimplifyScalarSlice( HloInstruction* slice) { // Only try to do this for effective scalars. We could do the same for slicing diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.h b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.h index 6ea0d835cef53a..7878708c3f5f0d 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.h +++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.h @@ -592,6 +592,8 @@ class AlgebraicSimplifierVisitor : public DfsHloRewriteVisitor { absl::Status HandleRsqrt(HloInstruction* rsqrt) override; + absl::Status HandleShuffle(HloInstruction* hlo) override; + absl::Status HandleSlice(HloInstruction* slice) override; absl::Status HandleSqrt(HloInstruction* sqrt) override; diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc index b439bf6ed0d202..bf5cc9ea7f9b67 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc @@ -15087,6 +15087,96 @@ TEST_F(AlgebraicSimplifierTest, DynamicSliceOfDynamicSlice) { m::ConstantScalar(6)))))); } +TEST_F(AlgebraicSimplifierTest, FusesShuffleRotates) { + constexpr absl::string_view kModuleStr = R"( + HloModule shuffle_module + + ENTRY main { + p0 = f32[5]{0} parameter(0) + r1 = f32[5]{0} shuffle(p0), dimensions={0}, mode=rotate, shifts={2} + ROOT r2 = f32[5]{0} shuffle(r1), dimensions={0}, mode=rotate, shifts={1} + } + )"; + ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_OK_AND_ASSIGN(bool changed, simplifier.Run(m.get())); + EXPECT_TRUE(changed); + + auto* root = m->entry_computation()->root_instruction(); + ASSERT_EQ(root->opcode(), HloOpcode::kShuffle); + auto* shuffle = Cast(root); + EXPECT_EQ(shuffle->operand(0)->opcode(), HloOpcode::kParameter); + ASSERT_EQ(shuffle->dimensions().size(), 1); + EXPECT_EQ(shuffle->dimensions()[0], 0); + EXPECT_EQ(shuffle->rotate().shifts(0), 3); +} + +TEST_F(AlgebraicSimplifierTest, RemovesShuffleRotateOnSplatBroadcast) { + constexpr absl::string_view kModuleStr = R"( + HloModule shuffle_module + + ENTRY main { + c = f32[] constant(1.0) + b = f32[5,5]{1,0} broadcast(c), dimensions={} + ROOT r = f32[5,5]{1,0} shuffle(b), dimensions={0}, mode=rotate, shifts={2} + } + )"; + ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_OK_AND_ASSIGN(bool changed, simplifier.Run(m.get())); + EXPECT_TRUE(changed); + + auto* root = m->entry_computation()->root_instruction(); + EXPECT_EQ(root->opcode(), HloOpcode::kBroadcast); +} + +TEST_F(AlgebraicSimplifierTest, RemovesNoopDimAndCanonicalizesShift) { + constexpr absl::string_view kModuleStr = R"( + HloModule shuffle_module + + ENTRY main { + p0 = f32[5,5]{1,0} parameter(0) + ROOT r = f32[5,5]{1,0} shuffle(p0), dimensions={1,0}, mode=rotate, + shifts={0,-1} + } + )"; + ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_OK_AND_ASSIGN(bool changed, simplifier.Run(m.get())); + EXPECT_TRUE(changed); + + auto* root = m->entry_computation()->root_instruction(); + ASSERT_EQ(root->opcode(), HloOpcode::kShuffle); + auto* shuffle = Cast(root); + ASSERT_EQ(shuffle->dimensions().size(), 1); + EXPECT_EQ(shuffle->dimensions()[0], 0); + EXPECT_EQ(shuffle->rotate().shifts(0), 4); +} + +TEST_F(AlgebraicSimplifierTest, CanonicalizesMultiDimShuffleRotate) { + constexpr absl::string_view kModuleStr = R"( + HloModule shuffle_module + + ENTRY main { + p0 = f32[5,6]{1,0} parameter(0) + ROOT r = f32[5,6]{1,0} shuffle(p0), dimensions={1,0}, mode=rotate, + shifts={3,-2} + } + )"; + ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_OK_AND_ASSIGN(bool changed, simplifier.Run(m.get())); + EXPECT_TRUE(changed); + + auto* root = m->entry_computation()->root_instruction(); + ASSERT_EQ(root->opcode(), HloOpcode::kShuffle); + auto* shuffle = Cast(root); + ASSERT_EQ(shuffle->dimensions().size(), 2); + EXPECT_EQ(shuffle->dimensions()[0], 0); + EXPECT_EQ(shuffle->rotate().shifts(0), 3); + EXPECT_EQ(shuffle->dimensions()[1], 1); + EXPECT_EQ(shuffle->rotate().shifts(1), 3); +} + } // namespace } // namespace xla -// Trivial comment to force rebuild diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc index 4f9f4e286bf1af..6cefdf8def86ab 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.cc @@ -19,6 +19,7 @@ limitations under the License. #include #include #include +#include #include #include @@ -47,7 +48,6 @@ limitations under the License. #include "xla/tsl/platform/errors.h" #include "xla/tsl/platform/statusor.h" #include "xla/xla_data.pb.h" -#include "tsl/platform/errors.h" namespace xla { @@ -447,6 +447,84 @@ absl::StatusOr PropagateIdenticalConstantArguments( return changed; } +// TODO(enver): Share CloneComputation with other passes, such as +// FlattenCallGraph, by moving to HloModule or HloComputation. +HloComputation* CloneComputation(HloModule* module, + HloComputation* computation) { + if (!module->has_schedule() || + !module->schedule().is_computation_scheduled(computation)) { + return module->AddEmbeddedComputation(computation->Clone()); + } + auto [clone, clone_sequence] = computation->CloneWithSchedule(); + HloComputation* clone_ptr = module->AddEmbeddedComputation(std::move(clone)); + module->schedule().set_sequence(clone_ptr, clone_sequence); + return clone_ptr; +} + +using SpecializationKey = HloConstantFolding::SpecializationKey; + +absl::StatusOr SpecializeCalls( + HloModule* module, HloComputation* computation, + absl::flat_hash_map& + specialization_cache, + std::vector& computation_versions) { + auto caller_instructions = computation->caller_instructions(); + if (caller_instructions.empty()) { + return false; + } + if (!absl::c_all_of(caller_instructions, [](const HloInstruction* instr) { + return instr->opcode() == HloOpcode::kCall; + })) { + return false; + } + + if (caller_instructions.size() > 1) { + // Sort the caller instructions by their unique id to make the compilation + // deterministic. + absl::c_sort(caller_instructions, + [](const HloInstruction* a, const HloInstruction* b) { + return a->unique_id() < b->unique_id(); + }); + } + + bool changed = false; + bool original_computation_used = false; + for (HloInstruction* caller : caller_instructions) { + std::vector> arguments; + arguments.reserve(caller->operand_count()); + for (const HloInstruction* operand : caller->operands()) { + if (operand->opcode() == HloOpcode::kConstant) { + arguments.push_back(LiteralSlice(operand->literal())); + } else { + arguments.push_back(std::nullopt); + } + } + SpecializationKey key{computation, std::move(arguments)}; + auto it = specialization_cache.find(key); + HloComputation* target_comp = nullptr; + if (it != specialization_cache.end()) { + target_comp = it->second; + } else { + if (!original_computation_used) { + // The first argument combination for this computation uses the original + // computation. + target_comp = computation; + original_computation_used = true; + } else { + target_comp = CloneComputation(module, computation); + computation_versions.push_back(target_comp); + } + specialization_cache.emplace(key, target_comp); + } + + if (caller->to_apply() != target_comp) { + caller->set_to_apply(target_comp); + changed = true; + } + } + return changed; +} + } // namespace absl::StatusOr HloConstantFolding::RunOnComputation( @@ -606,6 +684,7 @@ absl::StatusOr HloConstantFolding::RunOnComputation( absl::StatusOr HloConstantFolding::RunImpl( HloModule* module, const absl::flat_hash_set& execution_threads) { + specialization_cache_.clear(); // Limit the constant folding to 0 iterations to skip folding loops in the // default case. This retains the behavior from before while loop support in // HloEvaluator and may be revised. @@ -627,10 +706,21 @@ absl::StatusOr HloConstantFolding::RunImpl( // arguments from callers to callees. for (auto it = computations.rbegin(); it != computations.rend(); ++it) { HloComputation* computation = *it; - ABSL_ASSIGN_OR_RETURN(bool computation_changed, - RunOnComputation(computation, evaluator.get(), - is_foldable_computation)); - changed |= computation_changed; + // TODO(b/260601110): Early exit the computation if all callers are already + // constant folded. + std::vector computation_versions; + computation_versions.push_back(computation); + ABSL_ASSIGN_OR_RETURN(bool did_specialize, + SpecializeCalls(module, computation, specialization_cache_, + computation_versions)); + changed |= did_specialize; + + for (HloComputation* computation_version : computation_versions) { + ABSL_ASSIGN_OR_RETURN(bool version_changed, + RunOnComputation(computation_version, evaluator.get(), + is_foldable_computation)); + changed |= version_changed; + } } return changed; } diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.h b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.h index 2c4bb8a45049e7..00fdd843bdf336 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.h +++ b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding.h @@ -19,13 +19,16 @@ limitations under the License. #include #include #include +#include +#include #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" #include "absl/status/statusor.h" #include "absl/strings/string_view.h" -#include "xla/hlo/ir/hlo_module.h" +#include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/pass/hlo_pass_interface.h" +#include "xla/literal.h" #include "xla/shape.h" namespace xla { @@ -82,6 +85,49 @@ class HloConstantFolding : public HloModulePass { explicit HloConstantFolding(const Options& options) : options_(options) {} absl::string_view name() const override { return "constant_folding"; } + struct SpecializationKey { + const HloComputation* original_computation = nullptr; + std::vector> arguments; + + bool operator==(const SpecializationKey& other) const { + if (original_computation != other.original_computation || + arguments.size() != other.arguments.size()) { + return false; + } + for (size_t i = 0; i < arguments.size(); ++i) { + if (arguments[i].has_value() != other.arguments[i].has_value()) { + return false; + } + if (arguments[i].has_value()) { + if (!arguments[i]->Equal(*other.arguments[i], + /*layout_sensitive=*/true)) { + return false; + } + } + } + return true; + } + + template + friend H AbslHashValue(H h, const SpecializationKey& key) { + h = H::combine(std::move(h), key.original_computation, + key.arguments.size()); + for (const auto& arg : key.arguments) { + if (arg.has_value()) { + h = H::combine(std::move(h), true, Literal::AbslHashable(*arg)); + } else { + h = H::combine(std::move(h), false); + } + } + return h; + } + }; + + const absl::flat_hash_map& + specialization_cache() const { + return specialization_cache_; + } + protected: // Run constant folding operations on the given module. Returns whether the // module was changed (constant expressions folded). @@ -99,6 +145,7 @@ class HloConstantFolding : public HloModulePass { static std::atomic slow_op_counter_; Options options_; + absl::flat_hash_map specialization_cache_; }; } // namespace xla diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding_test.cc index 12f90085cdc2c1..22506ad1b04e88 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding_test.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_constant_folding_test.cc @@ -21,15 +21,18 @@ limitations under the License. #include #include +#include "absl/status/status_matchers.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/parser/hlo_parser.h" +#include "xla/hlo/testlib/filecheck.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" #include "xla/hlo/testlib/pattern_matcher_gmock.h" #include "xla/hlo/testlib/test.h" +#include "xla/hlo/transforms/simplifiers/flatten_call_graph.h" #include "xla/hlo/utils/hlo_matchers.h" #include "xla/layout_util.h" #include "xla/literal.h" @@ -46,6 +49,8 @@ limitations under the License. namespace xla { namespace { +using ::absl_testing::IsOkAndHolds; + namespace op = xla::testing::opcode_matchers; namespace m = xla::match; using HloConstantFoldingTest = HloHardwareIndependentTestBase; @@ -759,7 +764,7 @@ TEST_F(HloConstantFoldingTest, HloConstantFolding constant_folding; TF_ASSERT_OK_AND_ASSIGN(bool result, RunHloPass(&constant_folding, module.get())); - EXPECT_FALSE(result); + EXPECT_TRUE(result); } TEST_F(HloConstantFoldingTest, InterproceduralMultipleCallsitesSomeConstants) { @@ -792,11 +797,10 @@ TEST_F(HloConstantFoldingTest, InterproceduralMultipleCallsitesSomeConstants) { RunHloPass(&constant_folding, module.get())); EXPECT_TRUE(result); HloComputation* fn = module->GetComputationWithName("Fn"); - EXPECT_THAT(fn->root_instruction(), - GmockMatch(m::Add(m::Subtract(m::Multiply(m::ConstantScalar(1), - m::Parameter(1)), - m::Parameter(2)), - m::ConstantScalar(2)))); + EXPECT_THAT( + fn->root_instruction(), + GmockMatch(m::Add(m::Subtract(m::ConstantScalar(2), m::Parameter(2)), + m::ConstantScalar(2)))); } TEST_F(HloConstantFoldingTest, @@ -832,7 +836,7 @@ TEST_F(HloConstantFoldingTest, EXPECT_TRUE(result); HloComputation* fn = module->GetComputationWithName("Fn"); EXPECT_THAT(fn->root_instruction(), - GmockMatch(m::Subtract(m::Add(m::Add(), m::Parameter(1)), + GmockMatch(m::Subtract(m::Add(m::Constant(), m::Parameter(1)), m::Constant()))); } @@ -1400,5 +1404,441 @@ TEST_F(HloConstantFoldingTest, LateOptionsDontFoldGteWithControlDependency) { EXPECT_FALSE(result); } +TEST_F(HloConstantFoldingTest, ComputationCalledTwiceOnDifferentConstants) { + const std::string hlo_string = R"( + HloModule test + + // CHECK-LABEL: %subcomp ( + // CHECK: constant({15, 15}) + + // CHECK-LABEL: %subcomp.clone ( + // CHECK: constant({25, 25}) + subcomp { + p0 = s32[2] parameter(0) + p1 = s32[2] parameter(1) + c = s32[2] constant({5, 5}) + add0 = s32[2] add(p0, c) + ROOT add1 = s32[2] add(p1, add0) + } + + // CHECK-LABEL: ENTRY %entry ( + // CHECK: %call1 = s32[2]{0} call(%c10, %p0.1), to_apply=%subcomp + // CHECK: %call2 = s32[2]{0} call(%c20, %p1.1), to_apply=%subcomp.clone + ENTRY entry { + p0.1 = s32[2] parameter(0) + p1.1 = s32[2] parameter(1) + c10 = s32[2] constant({10, 10}) + c20 = s32[2] constant({20, 20}) + call1 = s32[2] call(c10, p0.1), to_apply=subcomp + call2 = s32[2] call(c20, p1.1), to_apply=subcomp + ROOT out = s32[2] add(call1, call2) + })"; + + { + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + } + + { + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + FlattenCallGraph flatten; + EXPECT_THAT(flatten.Run(module.get()), IsOkAndHolds(true)); + + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), + IsOkAndHolds(true)); + } +} + +TEST_F(HloConstantFoldingTest, ComputationCalledThriceOnTwoDifferentConstants) { + const std::string hlo_string = R"( + HloModule test + + // CHECK-LABEL: %subcomp ( + // CHECK: constant({15, 15}) + + // CHECK-LABEL: %subcomp.clone ( + // CHECK: constant({25, 25}) + + // CHECK-LABEL: %subcomp.clone.1 ( + // CHECK: constant({25, 25}) + subcomp { + p0 = s32[2] parameter(0) + p1 = s32[2] parameter(1) + c = s32[2] constant({5, 5}) + add0 = s32[2] add(p0, c) + ROOT add1 = s32[2] add(p1, add0) + } + + // CHECK-LABEL: ENTRY %entry ( + // CHECK: %call1 = s32[2]{0} call(%c10, %p0.1), to_apply=%subcomp + // CHECK: %call2 = s32[2]{0} call(%c20, %p0.1), to_apply=%subcomp.clone + // CHECK: %call3 = s32[2]{0} call(%c20, %p0.1), to_apply=%subcomp.clone.1 + ENTRY entry { + p0.1 = s32[2] parameter(0) + c10 = s32[2] constant({10, 10}) + c20 = s32[2] constant({20, 20}) + call1 = s32[2] call(c10, p0.1), to_apply=subcomp + call2 = s32[2] call(c20, p0.1), to_apply=subcomp + call3 = s32[2] call(c20, p0.1), to_apply=subcomp + add1 = s32[2] add(call1, call2) + ROOT out = s32[2] add(add1, call3) + })"; + + { + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + } + + { + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + FlattenCallGraph flatten; + EXPECT_THAT(flatten.Run(module.get()), IsOkAndHolds(true)); + + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), + IsOkAndHolds(true)); + } +} + +TEST_F(HloConstantFoldingTest, + ComputationCalledTwiceOnDifferentConstantsSingleArgument) { + const std::string hlo_string = R"( + HloModule test + + // CHECK-LABEL: %subcomp ( + // CHECK: constant({5, 5}) + subcomp { + p0 = s32[2] parameter(0) + c = s32[2] constant({5, 5}) + add0 = s32[2] add(p0, c) + ROOT add1 = s32[2] add(add0, add0) + } + + // CHECK-LABEL: ENTRY %entry ( + // CHECK: ROOT {{.*}} = s32[2]{0} constant({80, 80}) + ENTRY entry { + c10 = s32[2] constant({10, 10}) + c20 = s32[2] constant({20, 20}) + call1 = s32[2] call(c10), to_apply=subcomp + call2 = s32[2] call(c20), to_apply=subcomp + ROOT out = s32[2] add(call1, call2) + })"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), IsOkAndHolds(true)); +} + +TEST_F(HloConstantFoldingTest, + ComputationCalledTwiceOnConstantAndOnNonConstant) { + const std::string hlo_string = R"( + HloModule test + + // CHECK-LABEL: %subcomp ( + // CHECK: constant({15, 15}) + + // CHECK-LABEL: %subcomp.clone ( + // CHECK: constant({5, 5}) + subcomp { + p0 = s32[2] parameter(0) + p1 = s32[2] parameter(1) + c = s32[2] constant({5, 5}) + add0 = s32[2] add(p0, c) + ROOT add1 = s32[2] add(p1, add0) + } + + // CHECK-LABEL: ENTRY %entry ( + // CHECK: %call1 = s32[2]{0} call(%c10, %p0.1), to_apply=%subcomp + // CHECK: %call2 = s32[2]{0} call(%p0.1, %p1.1), to_apply=%subcomp.clone + // CHECK: ROOT %out = s32[2]{0} add(%call1, %call2) + ENTRY entry { + p0.1 = s32[2] parameter(0) + p1.1 = s32[2] parameter(1) + c10 = s32[2] constant({10, 10}) + call1 = s32[2] call(c10, p0.1), to_apply=subcomp + call2 = s32[2] call(p0.1, p1.1), to_apply=subcomp + ROOT out = s32[2] add(call1, call2) + })"; + + { + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + } + { + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + FlattenCallGraph flatten; + EXPECT_THAT(flatten.Run(module.get()), IsOkAndHolds(true)); + + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), + IsOkAndHolds(true)); + } +} + +TEST_F(HloConstantFoldingTest, + ComputationCalledTwiceOnConstantAndOnNonConstantSingleArgument) { + const std::string hlo_string = R"( + HloModule test + + // CHECK-LABEL: %subcomp ( + // CHECK: constant({5, 5}) + subcomp { + p0 = s32[2] parameter(0) + c = s32[2] constant({5, 5}) + add0 = s32[2] add(p0, c) + ROOT add1 = s32[2] add(add0, add0) + } + + // CHECK-LABEL: ENTRY %entry ( + // CHECK: %constant = s32[2]{0} constant({30, 30}) + // CHECK: %call2 = s32[2]{0} call(%p0.1), to_apply=%subcomp + // CHECK: ROOT %out = s32[2]{0} add(%constant, %call2) + ENTRY entry { + p0.1 = s32[2] parameter(0) + c10 = s32[2] constant({10, 10}) + call1 = s32[2] call(c10), to_apply=subcomp + call2 = s32[2] call(p0.1), to_apply=subcomp + ROOT out = s32[2] add(call1, call2) + })"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), IsOkAndHolds(true)); +} + +TEST_F(HloConstantFoldingTest, ComputationCalledTwiceOnSameConstants) { + const std::string hlo_string = R"( + HloModule test + + // CHECK-LABEL: %subcomp ( + // CHECK: constant({15, 15}) + subcomp { + p0 = s32[2] parameter(0) + p1 = s32[2] parameter(1) + c = s32[2] constant({5, 5}) + add0 = s32[2] add(p0, c) + ROOT add1 = s32[2] add(p1, add0) + } + + // CHECK-LABEL: ENTRY %entry ( + // CHECK: %call1 = s32[2]{0} call(%c10, %p0.1), to_apply=%subcomp + // CHECK: %call2 = s32[2]{0} call(%c10, %p1.1), to_apply=%subcomp + // CHECK: ROOT %out = s32[2]{0} add(%call1, %call2) + ENTRY entry { + p0.1 = s32[2] parameter(0) + p1.1 = s32[2] parameter(1) + c10 = s32[2] constant({10, 10}) + call1 = s32[2] call(c10, p0.1), to_apply=subcomp + call2 = s32[2] call(c10, p1.1), to_apply=subcomp + ROOT out = s32[2] add(call1, call2) + })"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), IsOkAndHolds(true)); +} + +TEST_F(HloConstantFoldingTest, + ComputationCalledTwiceOnSameConstantsSingleArgument) { + const std::string hlo_string = R"( + HloModule test + + // CHECK-LABEL: %subcomp ( + // CHECK: constant({5, 5}) + subcomp { + p0 = s32[2] parameter(0) + c = s32[2] constant({5, 5}) + add0 = s32[2] add(p0, c) + ROOT add1 = s32[2] add(add0, add0) + } + + // CHECK-LABEL: ENTRY %entry ( + // CHECK: ROOT {{.*}} = s32[2]{0} constant({60, 60}) + ENTRY entry { + p0.1 = s32[2] parameter(0) + c10 = s32[2] constant({10, 10}) + call1 = s32[2] call(c10), to_apply=subcomp + call2 = s32[2] call(c10), to_apply=subcomp + ROOT out = s32[2] add(call1, call2) + })"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), IsOkAndHolds(true)); +} + +TEST_F(HloConstantFoldingTest, ComputationCalledOnceOnConstants) { + const std::string hlo_string = R"( + HloModule test + + // CHECK-LABEL: %subcomp ( + // CHECK: constant({15, 15}) + subcomp { + p0 = s32[2] parameter(0) + p1 = s32[2] parameter(1) + c = s32[2] constant({5, 5}) + add0 = s32[2] add(p0, c) + ROOT add1 = s32[2] add(p1, add0) + } + + // CHECK-LABEL: ENTRY %entry ( + // CHECK: %call1 = s32[2]{0} call(%c10, %p0.1), to_apply=%subcomp + // CHECK: ROOT %out = s32[2]{0} negate(%call1) + ENTRY entry { + p0.1 = s32[2] parameter(0) + p1.1 = s32[2] parameter(1) + c10 = s32[2] constant({10, 10}) + call1 = s32[2] call(c10, p0.1), to_apply=subcomp + ROOT out = s32[2] negate(call1) + })"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), IsOkAndHolds(true)); +} + +TEST_F(HloConstantFoldingTest, NestedCallOnConstant) { + const std::string hlo_string = R"( + HloModule test + + // CHECK-LABEL: %bar ( + // CHECK-NEXT: %p0.bar = s32[2]{0} parameter(0) + // CHECK-NEXT: %constant.1 = s32[2]{0} constant({10, 10}) + // CHECK-NEXT: %p1.bar = s32[2]{0} parameter(1) + // CHECK-NEXT: ROOT %out.bar = s32[2]{0} add(%constant.1, %p1.bar) + bar { + p0.bar = s32[2] parameter(0) + p1.bar = s32[2] parameter(1) + negate.bar = s32[2] negate(p0.bar) + ROOT out.bar = s32[2] add(negate.bar, p1.bar) + } + + // CHECK-LABEL: %foo ( + // CHECK-NEXT: %p0.foo = s32[2]{0} parameter(0) + // CHECK-NEXT: %constant = s32[2]{0} constant({-10, -10}) + // CHECK-NEXT: %p1.foo = s32[2]{0} parameter(1) + // CHECK-NEXT: %call.bar = s32[2]{0} call(%constant, %p1.foo), to_apply=%bar + // CHECK-NEXT: ROOT %out.foo = s32[2]{0} negate(%call.bar) + foo { + p0.foo = s32[2] parameter(0) + p1.foo = s32[2] parameter(1) + negate.foo = s32[2] negate(p0.foo) + call.bar = s32[2] call(negate.foo, p1.foo), to_apply=bar + ROOT out.foo = s32[2] negate(call.bar) + } + + // CHECK-LABEL: ENTRY %entry ( + // CHECK-NEXT: %c10 = s32[2]{0} constant({10, 10}) + // CHECK-NEXT: %p0 = s32[2]{0} parameter(0) + // CHECK-NEXT: %call.foo = s32[2]{0} call(%c10, %p0), to_apply=%foo + // CHECK-NEXT: ROOT %out = s32[2]{0} negate(%call.foo) + ENTRY entry { + p0 = s32[2] parameter(0) + c10 = s32[2] constant({10, 10}) + call.foo = s32[2] call(c10, p0), to_apply=foo + ROOT out = s32[2] negate(call.foo) + })"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), IsOkAndHolds(true)); +} + +TEST_F(HloConstantFoldingTest, NestedCallAndFooCalledTwice) { + const std::string hlo_string = R"( + HloModule test + + // CHECK-LABEL: %bar ( + // CHECK-NEXT: %p0.bar = s32[2]{0} parameter(0) + // CHECK-NEXT: %negate.bar = s32[2]{0} negate(%p0.bar) + // CHECK-NEXT: ROOT %out.bar = s32[2]{0} add(%negate.bar, %negate.bar) + bar { + p0.bar = s32[2] parameter(0) + negate.bar = s32[2] negate(p0.bar) + ROOT out.bar = s32[2] add(negate.bar, negate.bar) + } + + // CHECK-LABEL: %foo ( + // CHECK-NEXT: %p0.foo = s32[2]{0} parameter(0) + // CHECK-NEXT: %constant.2 = s32[2]{0} constant({20, 20}) + // CHECK-NEXT: %p1.foo = s32[2]{0} parameter(1) + // CHECK-NEXT: ROOT %out.foo = s32[2]{0} add(%constant.2, %p1.foo) + foo { + p0.foo = s32[2] parameter(0) + p1.foo = s32[2] parameter(1) + negate.foo = s32[2] negate(p0.foo) + call.bar = s32[2] call(negate.foo), to_apply=bar + ROOT out.foo = s32[2] add(call.bar, p1.foo) + } + + // CHECK-LABEL: ENTRY %entry ( + // CHECK-NEXT: %c10 = s32[2]{0} constant({10, 10}) + // CHECK-NEXT: %p0 = s32[2]{0} parameter(0) + // CHECK-NEXT: %call.0 = s32[2]{0} call(%c10, %p0), to_apply=%foo + // CHECK-NEXT: %constant = s32[2]{0} constant({30, 30}) + // CHECK-NEXT: ROOT %out = s32[2]{0} add(%call.0, %constant) + ENTRY entry { + p0 = s32[2] parameter(0) + c10 = s32[2] constant({10, 10}) + call.0 = s32[2] call(c10, p0), to_apply=foo + call.1 = s32[2] call(c10, c10), to_apply=foo + ROOT out = s32[2] add(call.0, call.1) + })"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), IsOkAndHolds(true)); +} + +TEST_F(HloConstantFoldingTest, NestedCallWhereOriginalCallerIsFolded) { + const std::string hlo_string = R"( + HloModule test + + // CHECK-LABEL: %bar ( + // CHECK: %constant.1 = s32[2]{0} constant({3, 3}) + bar { + p0.bar = s32[2] parameter(0) + c1 = s32[2] constant({1, 1}) + c2 = s32[2] constant({2, 2}) + c3 = s32[2] add(c1, c2) + ROOT out.bar = s32[2] add(p0.bar, c3) + } + + foo { + p0.foo = s32[2] parameter(0) + p1.foo = s32[2] parameter(1) + call.bar = s32[2] call(p0.foo), to_apply=bar + ROOT out.foo = s32[2] add(call.bar, p1.foo) + } + + ENTRY entry { + p0 = s32[2] parameter(0) + p1 = s32[2] parameter(1) + c10 = s32[2] constant({10, 10}) + call.0 = s32[2] call(c10, p0), to_apply=foo + call.1 = s32[2] call(p1, p0), to_apply=foo + ROOT out = s32[2] add(call.0, call.1) + })"; + + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloConstantFolding constant_folding; + EXPECT_THAT(constant_folding.Run(module.get()), IsOkAndHolds(true)); + EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), IsOkAndHolds(true)); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/BUILD b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/BUILD index cdf5463c660cee..d6b5175a155620 100644 --- a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/BUILD +++ b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/BUILD @@ -129,8 +129,6 @@ cc_library( "//xla/mlir_hlo", "//xla/mlir_hlo:unregistered_attributes", "//xla/service:hlo_proto_cc", - "//xla/tsl/platform:errors", - "//xla/tsl/platform:statusor", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", diff --git a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc index b180a3292b3748..9890ec597b9e52 100644 --- a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc +++ b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc @@ -85,8 +85,6 @@ limitations under the License. #include "xla/service/hlo.pb.h" #include "xla/shape_util.h" #include "xla/status_macros.h" -#include "xla/tsl/platform/errors.h" -#include "xla/tsl/platform/statusor.h" #include "xla/util.h" #include "xla/xla_data.pb.h" @@ -1891,6 +1889,9 @@ absl::StatusOr HloFunctionImporter::ImportInstructionImpl( } return op.getOperation(); } + // TODO(b/565684884): Add conversion when Shuffle is supported in StableHLO. + case HloOpcode::kShuffle: + return InvalidArgument("Shuffle is not supported in StableHLO."); case HloOpcode::kRng: { auto shape = mlir::stablehlo::ConstantOp::create( *func_builder, loc, @@ -2141,8 +2142,8 @@ absl::StatusOr HloFunctionImporter::ImportInstructionImpl( &instruction->precision_config(), builder_))); // If the element types of the operands for convolution are different, - // insert a convert op to convert the operands to the common element type - // while preserving the values. + // insert a convert op to convert the operands to the common element + // type while preserving the values. auto lhs = operands[0]; auto rhs = operands[1]; auto lhs_element_type = instruction->operand(0)->shape().element_type(); @@ -2200,10 +2201,10 @@ absl::StatusOr HloFunctionImporter::ImportInstructionImpl( } case HloOpcode::kAdd: { - // HLO add ops on PRED elements are actually boolean or, but MHLO dialect - // AddOps on i1 are just addition with overflow; so, we have to implement - // the special behavior of HLO add ops on PRED here by creating an - // arith::OrIOp instead. + // HLO add ops on PRED elements are actually boolean or, but MHLO + // dialect AddOps on i1 are just addition with overflow; so, we have to + // implement the special behavior of HLO add ops on PRED here by + // creating an arith::OrIOp instead. if (instruction->shape().element_type() == PRED) { return mlir::stablehlo::OrOp::create(*func_builder, loc, result_type, operands, attributes) @@ -2227,8 +2228,8 @@ absl::StatusOr HloFunctionImporter::ImportInstructionImpl( } case HloOpcode::kConvert: { - // Convert to boolean is special, it requires a comparison to 0 instead of - // a truncation to i1, otherwise it is a 1-1 translation. + // Convert to boolean is special, it requires a comparison to 0 instead + // of a truncation to i1, otherwise it is a 1-1 translation. auto ranked_type = mlir::dyn_cast(result_type); mlir::IntegerType integer_type = (ranked_type) @@ -2279,10 +2280,10 @@ absl::StatusOr HloFunctionImporter::ImportInstructionImpl( "kind", mlir::mhlo::DomainKindAttr::get(func_builder->getContext(), *domain_kind))); - // In XLA, DomainMetadata is open-world, but in the proto, it is hardcoded - // to be ShardingMetadata. Thankfully, the only other implementation of - // DomainMetadata is OpName, which is generally used for debugging and - // never for compiling production models. + // In XLA, DomainMetadata is open-world, but in the proto, it is + // hardcoded to be ShardingMetadata. Thankfully, the only other + // implementation of DomainMetadata is OpName, which is generally used + // for debugging and never for compiling production models. // // Since this is hardcoded as such in the proto, we must follow suit. auto exit_metadata = ShardingMetadata::ToShardingMetadata( diff --git a/third_party/xla/xla/hlo/utils/hlo_matchers.h b/third_party/xla/xla/hlo/utils/hlo_matchers.h index c413dd80d45c7a..14b18c53d9bc14 100644 --- a/third_party/xla/xla/hlo/utils/hlo_matchers.h +++ b/third_party/xla/xla/hlo/utils/hlo_matchers.h @@ -387,6 +387,7 @@ HLO_MATCHER(SetDimensionSize); HLO_MATCHER(ShiftLeft); HLO_MATCHER(ShiftRightArithmetic); HLO_MATCHER(ShiftRightLogical); +HLO_MATCHER(Shuffle); HLO_MATCHER(Sign); HLO_MATCHER(Sinh); HLO_MATCHER(Slice); diff --git a/third_party/xla/xla/layout.h b/third_party/xla/xla/layout.h index 276243c2971f61..c25424c463abc2 100644 --- a/third_party/xla/xla/layout.h +++ b/third_party/xla/xla/layout.h @@ -404,6 +404,7 @@ class Layout { static constexpr int64_t kDefaultMemorySpace = 0; static constexpr int64_t kGenericFastMemorySpace = 1; + static constexpr int64_t kCollectiveMemorySpace = 7; static constexpr int64_t kHostMemorySpace = 5; static constexpr int64_t kUnpinnedHostMemorySpace = 6; int64_t memory_space() const { return memory_space_; } diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc index 6cf17ae96335fa..09432b11e5d080 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc @@ -955,6 +955,28 @@ LogicalResult TransposeOp::verify() { return success(); } +OpFoldResult TransposeOp::fold(FoldAdaptor adaptor) { + if (llvm::is_sorted(getPermutation())) { + return getVector(); + } + if (auto cst = dyn_cast_if_present(adaptor.getVector())) { + return cst.reshape(getType()); + } + if (auto prev_op = getVector().getDefiningOp()) { + ArrayRef prev_perm = prev_op.getPermutation(); + ArrayRef curr_perm = getPermutation(); + CHECK_EQ(prev_perm.size(), curr_perm.size()); + SmallVector composed_perm(curr_perm.size()); + for (size_t i = 0; i < curr_perm.size(); ++i) { + composed_perm[i] = prev_perm[curr_perm[i]]; + } + if (llvm::is_sorted(composed_perm)) { + return prev_op.getVector(); + } + } + return nullptr; +} + LogicalResult MemRefBitcastOp::verify() { auto src_ty = getInput().getType(); auto tgt_ty = getType(); diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td index e23c2715ba515f..1130919108db8c 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td @@ -1921,6 +1921,7 @@ def TPU_TransposeOp : TPU_Op<"transpose", [Pure]> { } }]; let hasVerifier = 1; + let hasFolder = 1; } def TPU_LogOp : TPU_Op<"log"> { diff --git a/third_party/xla/xla/pjrt/cpu/BUILD b/third_party/xla/xla/pjrt/cpu/BUILD index 8103e9d1f988c6..bb92234b2867c7 100644 --- a/third_party/xla/xla/pjrt/cpu/BUILD +++ b/third_party/xla/xla/pjrt/cpu/BUILD @@ -248,7 +248,6 @@ cc_library( "//xla/service/cpu:cpu_executable", "//xla/service/cpu:cpu_executable_run_options", "//xla/service/cpu:cpu_xfeed", - "//xla/service/cpu:executable_proto_cc", "//xla/service/llvm_ir:llvm_command_line_options", "//xla/stream_executor:device_address", "//xla/tsl/concurrency:async_value", diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.cc b/third_party/xla/xla/pjrt/cpu/cpu_client.cc index e30a0fe49f7877..0873c87b8560c9 100644 --- a/third_party/xla/xla/pjrt/cpu/cpu_client.cc +++ b/third_party/xla/xla/pjrt/cpu/cpu_client.cc @@ -71,6 +71,7 @@ limitations under the License. #include "xla/pjrt/cpu/cpu_async_execution_tracker.h" #include "xla/pjrt/cpu/cpu_device_memory.h" #include "xla/pjrt/cpu/cpu_event.h" +#include "xla/pjrt/cpu/execution_stream_event_map.h" #include "xla/pjrt/cpu/raw_buffer.h" #include "xla/pjrt/device_event.h" #include "xla/pjrt/device_event_utils.h" @@ -107,7 +108,6 @@ limitations under the License. #include "xla/service/cpu/cpu_executable.h" #include "xla/service/cpu/cpu_executable_run_options.h" #include "xla/service/cpu/cpu_xfeed.h" -#include "xla/service/cpu/executable.pb.h" #include "xla/service/device_assignment.h" #include "xla/service/dump.h" #include "xla/service/executable.h" @@ -129,7 +129,6 @@ limitations under the License. #include "xla/tsl/platform/threadpool.h" #include "xla/util.h" #include "xla/xla.pb.h" -#include "xla/xla_data.pb.h" #include "tsl/platform/denormal.h" #include "tsl/platform/fingerprint.h" #include "tsl/platform/setround.h" @@ -460,10 +459,15 @@ PjRtCpuRawClient::PjRtCpuRawClient( eigen_intraop_pool_->NumThreads()) : std::unique_ptr()), custom_intraop_device_(intra_op_device), + compile_thread_pool_(std::make_unique( + tsl::Env::Default(), GetThreadOptions(), "XLACompile", num_threads)), + execute_work_runner_(std::make_unique( + tsl::Env::Default(), "XLAExecute", num_threads, GetThreadOptions())), async_work_runner_(std::make_unique( - tsl::Env::Default(), "XLAPjRtCpuClient", num_threads)) {} + tsl::Env::Default(), "XLAPjRtCpuClient", num_threads, + GetThreadOptions())) {} -PjRtCpuRawClient::~PjRtCpuRawClient() {} +PjRtCpuRawClient::~PjRtCpuRawClient() = default; PjRtPluginAttributes GetDefaultCpuPluginAttributes() { PjRtPluginAttributes attrs; @@ -994,7 +998,7 @@ PjRtCpuRawClient::CompileInternal( params.layout_canonicalization_callback = std::move(layout_canonicalization_callback); params.num_threads = num_threads; - params.compile_thread_pool = async_work_runner()->thread_pool(); + params.compile_thread_pool = compile_thread_pool_.get(); params.aot_options = aot_options; params.process_index = process_index; params.collectives_exists = (collectives() != nullptr); @@ -1697,7 +1701,7 @@ PjRtRawLoadedExecutable::RawExecuteResult CpuPjRtRawLoadedExecutable::Execute( run_id_.ToInt(), std::move(ready_on_exit).Release()); PjRtDeviceEventSpan events_ref(input_deps); xla::ExecuteWhenReady( - events_ref, raw_client->async_work_runner(), + events_ref, raw_client->execute_work_runner(), [cpu_executable, buffer_alloc = std::move(buffer_alloc), buffer_alloc_and_copy = std::move(buffer_alloc_and_copy), execute_thunks = std::move(execute_thunks), diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.h b/third_party/xla/xla/pjrt/cpu/cpu_client.h index e249866e775ada..c467ebe46834be 100644 --- a/third_party/xla/xla/pjrt/cpu/cpu_client.h +++ b/third_party/xla/xla/pjrt/cpu/cpu_client.h @@ -122,6 +122,14 @@ class PjRtCpuRawClient : public PjRtRawClient { return async_work_runner_.get(); } + ThreadPoolAsyncWorkRunner* execute_work_runner() const { + return execute_work_runner_.get(); + } + + tsl::thread::ThreadPool* compile_thread_pool() const { + return compile_thread_pool_.get(); + } + tsl::thread::ThreadPool* eigen_intraop_pool() const { return eigen_intraop_pool_.get(); } @@ -292,6 +300,8 @@ class PjRtCpuRawClient : public PjRtRawClient { std::unique_ptr eigen_intraop_pool_; std::unique_ptr eigen_intraop_device_; const Eigen::ThreadPoolDevice* custom_intraop_device_ = nullptr; + std::unique_ptr compile_thread_pool_; + std::unique_ptr execute_work_runner_; std::unique_ptr async_work_runner_; }; 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 2feae674faa7b3..efcd52db06227a 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 @@ -3215,7 +3215,7 @@ constexpr char kAttrPlacedModule[] = R"( s = f32[512,1024] custom-call(t), custom_call_target="RecordBufferAddress", api_version=API_VERSION_TYPED_FFI, output_to_operand_aliasing={{}: (0, {})}, - frontend_attributes={results_memory_spaces="{0:1}"} + frontend_attributes={results_memory_spaces="{0:7}"} ROOT r = f32[512,1024] add(s, s) })"; diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.cc index a4cb5769b865ea..b25d2b2f6a63f9 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.cc +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_topology_description.cc @@ -425,6 +425,7 @@ StreamExecutorGpuTopologyDescription::GetMemorySpaceKindForShape( return GetMemorySpaceKindIds()[1]; break; case Layout::kGenericFastMemorySpace: + case Layout::kCollectiveMemorySpace: case Layout::kDefaultMemorySpace: break; default: 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 2d5ad0a0168708..362c40dfb1a82c 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 @@ -1761,6 +1761,7 @@ absl::StatusOr MemoryKindFromSimpleShape( case Layout::kHostMemorySpace: return PinnedHostMemorySpace::kKind; case Layout::kGenericFastMemorySpace: + case Layout::kCollectiveMemorySpace: case Layout::kDefaultMemorySpace: return default_memory_kind; default: diff --git a/third_party/xla/xla/pjrt/se/stream_executor_executable.cc b/third_party/xla/xla/pjrt/se/stream_executor_executable.cc index ace503c89b6b52..9c80957bd98a10 100644 --- a/third_party/xla/xla/pjrt/se/stream_executor_executable.cc +++ b/third_party/xla/xla/pjrt/se/stream_executor_executable.cc @@ -262,6 +262,7 @@ absl::StatusOr MemoryKindFromLayout( return PinnedHostMemorySpace::kKind; case Layout::kGenericFastMemorySpace: case Layout::kDefaultMemorySpace: + case Layout::kCollectiveMemorySpace: return default_memory_kind; default: return InvalidArgument("Unexpected memory space %d in output layout", diff --git a/third_party/xla/xla/python/_hlo.pyi b/third_party/xla/xla/python/_hlo.pyi index f208055c5ae3ba..5b1505f46a492e 100644 --- a/third_party/xla/xla/python/_hlo.pyi +++ b/third_party/xla/xla/python/_hlo.pyi @@ -321,37 +321,39 @@ class HloOpcode(enum.Enum): kShiftRightLogical = 117 - kSign = 118 + kShuffle = 118 - kSin = 119 + kSign = 119 - kSinh = 120 + kSin = 120 - kSlice = 121 + kSinh = 121 - kSort = 122 + kSlice = 122 - kSqrt = 123 + kSort = 123 - kStochasticConvert = 124 + kSqrt = 124 - kSubtract = 125 + kStochasticConvert = 125 - kTan = 126 + kSubtract = 126 - kTanh = 127 + kTan = 127 - kTopK = 128 + kTanh = 128 - kTranspose = 129 + kTopK = 129 - kTriangularSolve = 130 + kTranspose = 130 - kTuple = 131 + kTriangularSolve = 131 - kWhile = 132 + kTuple = 132 - kXor = 133 + kWhile = 133 + + kXor = 134 class Layout: @overload diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD index bfb0cec85ccb07..45d7e446bcae49 100644 --- a/third_party/xla/xla/service/BUILD +++ b/third_party/xla/xla/service/BUILD @@ -551,6 +551,7 @@ xla_cc_test( ":hlo_proto_cc", ":shape_inference", "//xla:shape_util", + "//xla:shuffle", "//xla:xla_data_proto_cc", "//xla/hlo/builder:padding", "//xla/hlo/ir:hlo", @@ -561,6 +562,7 @@ xla_cc_test( "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", + "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", @@ -818,6 +820,7 @@ xla_cc_test( "//xla:comparison_util", "//xla:literal_util", "//xla:shape_util", + "//xla:shuffle", "//xla:util", "//xla:window_util", "//xla:xla_data_proto_cc", @@ -1050,6 +1053,7 @@ xla_cc_test( srcs = ["platform_util_test.cc"], deps = [ ":platform_util", + "//xla/stream_executor:mock_platform", "//xla/stream_executor/cuda:cuda_platform_id", "//xla/stream_executor/host:host_platform_id", "//xla/stream_executor/rocm:rocm_platform_id", @@ -1631,7 +1635,7 @@ cc_library( visibility = internal_visibility([ "//xla/internal:hwi_internal", "//xla/backends/gpu/ffi:__pkg__", - "//xla/backends/gpu/libraries/native_custom_call_thunks:handler_allowlist", + "//xla:native_custom_call_handler_allowlist", "//xla/backends/gpu/runtime:__subpackages__", "//xla/backends/gpu/tests:__pkg__", "//xla/backends/gpu/transforms:__subpackages__", @@ -1890,7 +1894,7 @@ cc_library( ], visibility = internal_visibility([ ":friends", - "//xla/backends/gpu/libraries/native_custom_call_thunks:handler_allowlist", + "//xla:native_custom_call_handler_allowlist", ]), deps = [ ":buffer_assignment_proto_cc", @@ -4701,6 +4705,7 @@ xla_cc_test( srcs = ["while_util_test.cc"], deps = [ ":while_util", + "//xla:shape_util", "//xla:util", "//xla/hlo/ir:hlo", "//xla/hlo/testlib:hlo_hardware_independent_test_base", diff --git a/third_party/xla/xla/service/cpu/cpu_instruction_fusion.cc b/third_party/xla/xla/service/cpu/cpu_instruction_fusion.cc index d254d0ea4e144f..c34b0af2b71fd8 100644 --- a/third_party/xla/xla/service/cpu/cpu_instruction_fusion.cc +++ b/third_party/xla/xla/service/cpu/cpu_instruction_fusion.cc @@ -221,6 +221,7 @@ bool CpuInstructionFusion::IsExpensive(const HloInstruction& instruction) { case HloOpcode::kShiftLeft: case HloOpcode::kShiftRightArithmetic: case HloOpcode::kShiftRightLogical: + case HloOpcode::kShuffle: case HloOpcode::kSlice: case HloOpcode::kStochasticConvert: case HloOpcode::kSubtract: diff --git a/third_party/xla/xla/service/float_support.cc b/third_party/xla/xla/service/float_support.cc index 59a6be3ea620f7..c1854a183e49fe 100644 --- a/third_party/xla/xla/service/float_support.cc +++ b/third_party/xla/xla/service/float_support.cc @@ -133,6 +133,7 @@ bool FloatSupport::EffectiveOperandPrecisionIsOutputPrecision( case HloOpcode::kRaggedAllToAll: case HloOpcode::kReshape: case HloOpcode::kReverse: + case HloOpcode::kShuffle: case HloOpcode::kSlice: case HloOpcode::kSort: case HloOpcode::kTranspose: diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD index 8fdb54cd4a0841..394417048f9d21 100644 --- a/third_party/xla/xla/service/gpu/BUILD +++ b/third_party/xla/xla/service/gpu/BUILD @@ -2428,9 +2428,7 @@ xla_test( deps = [ "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", - "//xla/hlo/testlib:pattern_matcher_gmock", "//xla/service:hlo_module_config", - "//xla/service:pattern_matcher", "//xla/tests:xla_internal_test_main", "//xla/tests/restricted:hlo_test_base_legacy", "@com_google_absl//absl/log", diff --git a/third_party/xla/xla/service/gpu/auto_sharding_gpu_compiler_test.cc b/third_party/xla/xla/service/gpu/auto_sharding_gpu_compiler_test.cc index 512bdd512ec446..d689536b66936b 100644 --- a/third_party/xla/xla/service/gpu/auto_sharding_gpu_compiler_test.cc +++ b/third_party/xla/xla/service/gpu/auto_sharding_gpu_compiler_test.cc @@ -16,14 +16,12 @@ limitations under the License. #include #include -#include #include #include "absl/log/log.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_module.h" -#include "xla/hlo/testlib/pattern_matcher_gmock.h" +#include "xla/hlo/ir/hlo_sharding.h" #include "xla/service/hlo_module_config.h" -#include "xla/service/pattern_matcher.h" #include "xla/tests/restricted/hlo_test_base_legacy.h" #include "xla/xla_data.pb.h" #include "tsl/platform/logging.h" @@ -32,8 +30,6 @@ namespace xla { namespace gpu { namespace { -namespace m = ::xla::match; - class AutoShardingTest : public HloTestBaseLegacy { protected: const char* const dot_hlo_string_ = R"( @@ -68,7 +64,10 @@ TEST_F(AutoShardingTest, MatMulWithoutAutosharding) { auto* instruction = compiled_module->entry_computation()->parameter_instruction(0); VLOG(2) << instruction->ToString(); - EXPECT_THAT(instruction, GmockMatch(m::Op().WithSharding("{replicated}"))); + EXPECT_FALSE(instruction->has_sharding()); + ASSERT_TRUE(compiled_module->has_spmd_parameters_shardings()); + EXPECT_EQ(compiled_module->spmd_parameters_shardings()[0], + HloSharding::Replicate()); } } // namespace diff --git a/third_party/xla/xla/service/gpu/execution_stream_assignment_test.cc b/third_party/xla/xla/service/gpu/execution_stream_assignment_test.cc index d50dd053f74420..76826c35af37c4 100644 --- a/third_party/xla/xla/service/gpu/execution_stream_assignment_test.cc +++ b/third_party/xla/xla/service/gpu/execution_stream_assignment_test.cc @@ -109,8 +109,8 @@ TEST_F(ExecutionStreamAssignmentTest, CopyStartStreamIdTest) { HloModule Module ENTRY CopyStartAndCopyDone { - p0 = f32[2,3]{1,0:S(1)} parameter(0) - copy-start = (f32[2,3]{1,0:S(2)}, f32[2,3]{1,0:S(1)}, u32[]) copy-start(p0) + p0 = f32[2,3]{1,0:S(7)} parameter(0) + copy-start = (f32[2,3]{1,0:S(2)}, f32[2,3]{1,0:S(7)}, u32[]) copy-start(p0) ROOT copy-done = f32[2,3]{1,0:S(2)} copy-done(copy-start) } )"; diff --git a/third_party/xla/xla/service/gpu/gpu_compiler_test.cc b/third_party/xla/xla/service/gpu/gpu_compiler_test.cc index 0283257a50c84e..d11dc58f45c9cf 100644 --- a/third_party/xla/xla/service/gpu/gpu_compiler_test.cc +++ b/third_party/xla/xla/service/gpu/gpu_compiler_test.cc @@ -1909,7 +1909,7 @@ TEST_F(GpuCompilerTest, MosaicMultimemRequiresSymmetricMemoryCopies) { p_multimem = s32[1] parameter(0) p_non_coll = s32[1] parameter(1) - cc_multimem = (s32[1]{0}) custom-call(p_multimem), custom_call_target="mosaic_gpu_v2", frontend_attributes={operands_memory_spaces="{0:1}",results_memory_spaces="{0:1}"}, api_version=API_VERSION_TYPED_FFI + cc_multimem = (s32[1]{0}) custom-call(p_multimem), custom_call_target="mosaic_gpu_v2", frontend_attributes={operands_memory_spaces="{0:7}",results_memory_spaces="{0:7}"}, api_version=API_VERSION_TYPED_FFI res_multimem = s32[1] get-tuple-element(cc_multimem), index=0 cc_non_coll = (s32[1]{0}) custom-call(p_non_coll), custom_call_target="mosaic_gpu_v2", api_version=API_VERSION_TYPED_FFI @@ -1931,16 +1931,16 @@ TEST_F(GpuCompilerTest, MosaicMultimemRequiresSymmetricMemoryCopies) { // CHECK-DAG: [[P_NON:%[^ ]+]] = s32[1]{0} parameter(1) // Multimem input parameters are copied to S1 - // CHECK-DAG: [[COPY_MULTI_IN:%copy[^ ]*]] = s32[1]{0:S(1)} copy([[P_MULTI]]) + // CHECK-DAG: [[COPY_MULTI_IN:%copy[^ ]*]] = s32[1]{0:S(7)} copy([[P_MULTI]]) - // CHECK-DAG: [[CC_MULTI:%[^ ]+]] = (s32[1]{0:S(1)}) custom-call([[COPY_MULTI_IN]]){{.*}}results_memory_spaces={0:1} + // CHECK-DAG: [[CC_MULTI:%[^ ]+]] = (s32[1]{0:S(7)}) custom-call([[COPY_MULTI_IN]]){{.*}}results_memory_spaces={0:7} // CHECK-DAG: [[CC_NON:%[^ ]+]] = (s32[1]{0}) custom-call([[P_NON]]) // Extracting from the 1-element tuples returned by custom calls (all index=0) - // CHECK-DAG: [[GTE_MULTI:%[^ ]+]] = s32[1]{0:S(1)} get-tuple-element([[CC_MULTI]]), index=0 + // CHECK-DAG: [[GTE_MULTI:%[^ ]+]] = s32[1]{0:S(7)} get-tuple-element([[CC_MULTI]]), index=0 // CHECK-DAG: [[GTE_NON:%[^ ]+]] = s32[1]{0} get-tuple-element([[CC_NON]]), index=0 - // The S1 -> S0 isolation copies + // The S7 -> S0 isolation copies // CHECK-DAG: [[COPY_OUT_MULTI:%copy[^ ]*]] = s32[1]{0} copy([[GTE_MULTI]]) // The final ROOT is pure S0 @@ -1972,7 +1972,7 @@ TEST_F(GpuCompilerTest, MosaicCollectiveMetadataRequiresSymmetricMemoryCopies) { HloModule test ENTRY main { p = s32[1] parameter(0) - mosaic = (s32[1]{0}) custom-call(p), custom_call_target="mosaic_gpu_v2", api_version=API_VERSION_TYPED_FFI, backend_config={uses_xla_collective_metadata=true}, frontend_attributes={operands_memory_spaces="{0:1}", results_memory_spaces="{0:1}"} + mosaic = (s32[1]{0}) custom-call(p), custom_call_target="mosaic_gpu_v2", api_version=API_VERSION_TYPED_FFI, backend_config={uses_xla_collective_metadata=true}, frontend_attributes={operands_memory_spaces="{0:7}", results_memory_spaces="{0:7}"} ROOT res = s32[1]{0} get-tuple-element(mosaic), index=0 } )"; @@ -1993,9 +1993,9 @@ TEST_F(GpuCompilerTest, MosaicCollectiveMetadataRequiresSymmetricMemoryCopies) { constexpr absl::string_view kExpected = R"( // CHECK: [[P:%[^ ]+]] = s32[1]{0} parameter(0) - // CHECK: [[COPY_IN:%copy[^ ]*]] = s32[1]{0:S(1)} copy([[P]]) - // CHECK: [[MOSAIC:%[^ ]+]] = (s32[1]{0:S(1)}) custom-call([[COPY_IN]]){{.*}}backend_config={uses_xla_collective_metadata=true} - // CHECK: [[GTE:%[^ ]+]] = s32[1]{0:S(1)} get-tuple-element([[MOSAIC]]), index=0 + // CHECK: [[COPY_IN:%copy[^ ]*]] = s32[1]{0:S(7)} copy([[P]]) + // CHECK: [[MOSAIC:%[^ ]+]] = (s32[1]{0:S(7)}) custom-call([[COPY_IN]]){{.*}}backend_config={uses_xla_collective_metadata=true} + // CHECK: [[GTE:%[^ ]+]] = s32[1]{0:S(7)} get-tuple-element([[MOSAIC]]), index=0 // CHECK: ROOT [[COPY_OUT:%copy[^ ]*]] = s32[1]{0} copy([[GTE]]) )"; EXPECT_THAT( @@ -2064,8 +2064,8 @@ TEST_P(OneShotRaggedAllToAllMemSpaceTest, DirectUsage) { constexpr absl::string_view kS1TwoCopies = R"( // CHECK: %output = f32[16]{0} parameter(1) - // CHECK: [[COPY1:%copy[0-9.]*]] = f32[16]{0:S(1)} copy(%output) - // CHECK: [[RA2A:%[^ ]+]] = f32[16]{0:S(1)} ragged-all-to-all(%input, [[COPY1]], + // CHECK: [[COPY1:%copy[0-9.]*]] = f32[16]{0:S(7)} copy(%output) + // CHECK: [[RA2A:%[^ ]+]] = f32[16]{0:S(7)} ragged-all-to-all(%input, [[COPY1]], // CHECK: ROOT %copy.{{[0-9]+}} = f32[16]{0} copy([[RA2A]]) )"; @@ -2140,10 +2140,10 @@ TEST_P(OneShotRaggedAllToAllMemSpaceTest, LoopUsage) { constexpr absl::string_view kS1TwoCopies = R"( // CHECK: [[OUTPUT1:%output[0-9.]*]] = f32[16]{0} parameter(1) - // CHECK: [[COPY1:%copy[0-9.]*]] = f32[16]{0:S(1)} copy([[OUTPUT1]]) - // CHECK: %tuple = (s32[], f32[16]{0}, f32[16]{0:S(1)}) tuple(%copy{{.*}}, %input{{.*}}, [[COPY1]]) - // CHECK: %while = (s32[], f32[16]{0}, f32[16]{0:S(1)}) while(%tuple) - // CHECK: [[RESULT:%result[0-9.]*]] = f32[16]{0:S(1)} get-tuple-element(%while), index=2 + // CHECK: [[COPY1:%copy[0-9.]*]] = f32[16]{0:S(7)} copy([[OUTPUT1]]) + // CHECK: %tuple = (s32[], f32[16]{0}, f32[16]{0:S(7)}) tuple(%copy{{.*}}, %input{{.*}}, [[COPY1]]) + // CHECK: %while = (s32[], f32[16]{0}, f32[16]{0:S(7)}) while(%tuple) + // CHECK: [[RESULT:%result[0-9.]*]] = f32[16]{0:S(7)} get-tuple-element(%while), index=2 // CHECK: ROOT %copy{{.*}} = f32[16]{0} copy([[RESULT]]) )"; @@ -2212,8 +2212,8 @@ ENTRY test_computation { )"; constexpr absl::string_view kS1TwoCopies = R"( - // CHECK: [[COPY0:%copy[0-9.]*]] = u32[2]{0:S(1)} copy(%p) - // CHECK: [[PERMUTE:%[^ ]+]] = u32[2]{0:S(1)} collective-permute([[COPY0]]) + // CHECK: [[COPY0:%copy[0-9.]*]] = u32[2]{0:S(7)} copy(%p) + // CHECK: [[PERMUTE:%[^ ]+]] = u32[2]{0:S(7)} collective-permute([[COPY0]]) // CHECK: ROOT %copy{{.*}} = u32[2]{0} copy([[PERMUTE]]) )"; @@ -2314,10 +2314,10 @@ TEST_P(RequiresCollectiveSymmetricMemorySpaceTest, LoopUsage) { constexpr absl::string_view kS1TwoCopies = R"( // CHECK:ENTRY %entry_computation (input: u32[2]) -> u32[2] { // CHECK: %input = u32[2]{0} parameter(0) - // CHECK: [[COPY1:%copy[0-9.]*]] = u32[2]{0:S(1)} copy(%input) - // CHECK: %tuple = (s32[], u32[2]{0:S(1)}) tuple(%copy{{.*}}, [[COPY1]]) - // CHECK: %while = (s32[], u32[2]{0:S(1)}) while(%tuple) - // CHECK: [[RESULT:%result[0-9.]*]] = u32[2]{0:S(1)} get-tuple-element(%while) + // CHECK: [[COPY1:%copy[0-9.]*]] = u32[2]{0:S(7)} copy(%input) + // CHECK: %tuple = (s32[], u32[2]{0:S(7)}) tuple(%copy{{.*}}, [[COPY1]]) + // CHECK: %while = (s32[], u32[2]{0:S(7)}) while(%tuple) + // CHECK: [[RESULT:%result[0-9.]*]] = u32[2]{0:S(7)} get-tuple-element(%while) // CHECK: ROOT %copy{{.*}} = u32[2]{0} copy([[RESULT]]) )"; @@ -2428,8 +2428,8 @@ TEST_P(GpuCompilerParametersCopyCollectiveMemoryTest, DirectUsage) { )"; constexpr absl::string_view kS1TwoCopies = R"( - // CHECK: [[COPY_IN:%copy[0-9.]*]] = s32[1]{0:S(1)} copy(%parameter_used_by_collective) - // CHECK: [[ALL_REDUCE:%[^ ]+]] = s32[1]{0:S(1)} all-reduce([[COPY_IN]]) + // CHECK: [[COPY_IN:%copy[0-9.]*]] = s32[1]{0:S(7)} copy(%parameter_used_by_collective) + // CHECK: [[ALL_REDUCE:%[^ ]+]] = s32[1]{0:S(7)} all-reduce([[COPY_IN]]) // CHECK: ROOT %copy.{{[0-9]+}} = s32[1]{0} copy([[ALL_REDUCE]]) )"; @@ -2538,10 +2538,10 @@ TEST_P(GpuCompilerParametersCopyCollectiveMemoryTest, LoopUsage) { constexpr absl::string_view kS1TwoCopies = R"( // CHECK: %input = s32[1]{0} parameter(0) - // CHECK: %copy.{{[0-9]+}} = s32[1]{0:S(1)} copy(%input) - // CHECK: %tuple = (s32[], s32[1]{0:S(1)}) tuple({{.*}}, %copy.{{[0-9]+}}) - // CHECK: %while = (s32[], s32[1]{0:S(1)}) while(%tuple) - // CHECK: %result.2.0 = s32[1]{0:S(1)} get-tuple-element(%while), index=1 + // CHECK: %copy.{{[0-9]+}} = s32[1]{0:S(7)} copy(%input) + // CHECK: %tuple = (s32[], s32[1]{0:S(7)}) tuple({{.*}}, %copy.{{[0-9]+}}) + // CHECK: %while = (s32[], s32[1]{0:S(7)}) while(%tuple) + // CHECK: %result.2.0 = s32[1]{0:S(7)} get-tuple-element(%while), index=1 // CHECK: ROOT %copy.{{[0-9]+}} = s32[1]{0} copy(%result.2.0) )"; @@ -2754,8 +2754,8 @@ TEST_P(FrontendAttributesMemorySpaceTest, DirectUsage) { custom_call_target="__xla_test_mock_custom_call_f32", api_version=API_VERSION_TYPED_FFI, frontend_attributes={ - operands_memory_spaces="{0:1}", - results_memory_spaces="{0:1}" + operands_memory_spaces="{0:7}", + results_memory_spaces="{0:7}" } } )"; @@ -2779,8 +2779,8 @@ TEST_P(FrontendAttributesMemorySpaceTest, DirectUsage) { // Regardless of the input_output_alias it should be two copies constexpr absl::string_view expected_check = R"( // CHECK: %p = f32[16]{0} parameter(0) - // CHECK: [[COPY0:%copy[0-9.]*]] = f32[16]{0:S(1)} copy(%p) - // CHECK: %cc = f32[16]{0:S(1)} custom-call([[COPY0]]) + // CHECK: [[COPY0:%copy[0-9.]*]] = f32[16]{0:S(7)} copy(%p) + // CHECK: %cc = f32[16]{0:S(7)} custom-call([[COPY0]]) // CHECK: ROOT %copy{{.*}} = f32[16]{0} copy(%cc) )"; @@ -2811,8 +2811,8 @@ TEST_P(FrontendAttributesMemorySpaceTest, LoopUsage) { custom_call_target="__xla_test_mock_custom_call_f32", api_version=API_VERSION_TYPED_FFI, frontend_attributes={ - operands_memory_spaces="{0:1}", - results_memory_spaces="{0:1}" + operands_memory_spaces="{0:7}", + results_memory_spaces="{0:7}" } new_loop_counter = s32[] add(loop_counter, s32[] constant(1)) ROOT result = (s32[], f32[16]) tuple(new_loop_counter, cc) @@ -2846,10 +2846,10 @@ TEST_P(FrontendAttributesMemorySpaceTest, LoopUsage) { // Regardless of the input_output_alias it should be two copies constexpr absl::string_view expected_check = R"( // CHECK: %input = f32[16]{0} parameter(0) - // CHECK: [[COPY0:%copy[0-9.]*]] = f32[16]{0:S(1)} copy(%input) - // CHECK: %tuple = (s32[], f32[16]{0:S(1)}) tuple(%copy{{.*}}, [[COPY0]]) - // CHECK: %while = (s32[], f32[16]{0:S(1)}) while(%tuple) - // CHECK: [[RESULT:%result[0-9.]*]] = f32[16]{0:S(1)} get-tuple-element(%while), index=1 + // CHECK: [[COPY0:%copy[0-9.]*]] = f32[16]{0:S(7)} copy(%input) + // CHECK: %tuple = (s32[], f32[16]{0:S(7)}) tuple(%copy{{.*}}, [[COPY0]]) + // CHECK: %while = (s32[], f32[16]{0:S(7)}) while(%tuple) + // CHECK: [[RESULT:%result[0-9.]*]] = f32[16]{0:S(7)} get-tuple-element(%while), index=1 // CHECK: ROOT %copy{{.*}} = f32[16]{0} copy([[RESULT]]) )"; @@ -2867,13 +2867,13 @@ TEST_F(GpuCompilerTest, HloModule test, input_output_alias={ {}: (0, {}) } ENTRY test_computation { - p = f32[16]{0:S(1)} parameter(0) - ROOT cc = f32[16]{0:S(1)} custom-call(p), + p = f32[16]{0:S(7)} parameter(0) + ROOT cc = f32[16]{0:S(7)} custom-call(p), custom_call_target="__xla_test_mock_custom_call_f32", api_version=API_VERSION_TYPED_FFI, frontend_attributes={ - operands_memory_spaces="{0:1}", - results_memory_spaces="{0:1}" + operands_memory_spaces="{0:7}", + results_memory_spaces="{0:7}" } } )"; @@ -2923,13 +2923,13 @@ TEST_F(GpuCompilerTest, HloModule test, input_output_alias={ {}: (0, {}) } ENTRY test_computation { - p = f32[16]{0:S(1)} parameter(0) - ROOT cc = f32[16]{0:S(1)} custom-call(p), + p = f32[16]{0:S(7)} parameter(0) + ROOT cc = f32[16]{0:S(7)} custom-call(p), custom_call_target="__xla_test_mock_custom_call_f32", api_version=API_VERSION_TYPED_FFI, frontend_attributes={ - operands_memory_spaces="{0:1}", - results_memory_spaces="{0:1}" + operands_memory_spaces="{0:7}", + results_memory_spaces="{0:7}" } } )"; @@ -2963,13 +2963,13 @@ TEST_F(GpuCompilerTest, HloModule test ENTRY test_computation { - p = f32[16]{0:S(1)} parameter(0) - ROOT cc = f32[16]{0:S(1)} custom-call(p), + p = f32[16]{0:S(7)} parameter(0) + ROOT cc = f32[16]{0:S(7)} custom-call(p), custom_call_target="__xla_test_mock_custom_call_f32", api_version=API_VERSION_TYPED_FFI, frontend_attributes={ - operands_memory_spaces="{0:1}", - results_memory_spaces="{0:1}" + operands_memory_spaces="{0:7}", + results_memory_spaces="{0:7}" } } )"; @@ -3004,8 +3004,8 @@ TEST_F(GpuCompilerTest, custom_call_target="__xla_test_mock_custom_call_f32", api_version=API_VERSION_TYPED_FFI, frontend_attributes={ - operands_memory_spaces="{0:1}", - results_memory_spaces="{0:1}" + operands_memory_spaces="{0:7}", + results_memory_spaces="{0:7}" } } )"; @@ -3040,8 +3040,8 @@ TEST_F(GpuCompilerTest, custom_call_target="__xla_test_mock_custom_call_f32", api_version=API_VERSION_TYPED_FFI, frontend_attributes={ - operands_memory_spaces="{0:1}", - results_memory_spaces="{0:1}" + operands_memory_spaces="{0:7}", + results_memory_spaces="{0:7}" } } )"; @@ -3071,16 +3071,16 @@ TEST_F(GpuCompilerTest, HloModule test, input_output_alias={ {0}: (0, {}) } ENTRY test_computation { - p0 = f32[16]{0:S(1)} parameter(0) - cc = (f32[16]{0:S(1)}, f32[16]{0:S(1)}) custom-call(p0), + p0 = f32[16]{0:S(7)} parameter(0) + cc = (f32[16]{0:S(7)}, f32[16]{0:S(7)}) custom-call(p0), custom_call_target="mock_call", frontend_attributes={ - operands_memory_spaces="{0:1}", - results_memory_spaces="{0:1,1:1}" + operands_memory_spaces="{0:7}", + results_memory_spaces="{0:7,1:7}" } - elem0 = f32[16]{0:S(1)} get-tuple-element(cc), index=0 - elem1 = f32[16]{0:S(1)} get-tuple-element(cc), index=1 - ROOT out = (f32[16]{0:S(1)}, f32[16]{0:S(1)}) tuple(elem1, elem0) + elem0 = f32[16]{0:S(7)} get-tuple-element(cc), index=0 + elem1 = f32[16]{0:S(7)} get-tuple-element(cc), index=1 + ROOT out = (f32[16]{0:S(7)}, f32[16]{0:S(7)}) tuple(elem1, elem0) } )"; @@ -3243,7 +3243,7 @@ TEST_F(GpuCompilerTest, SymmetricBuffersFilter) { // - channel 3 does NOT have S(1) (f32, 8192 bytes - filtered out by size) constexpr absl::string_view expected_check = R"( - // CHECK-DAG: f32[1024]{0:S(1)}{{.*}}all-reduce{{.*}}channel_id=1 + // CHECK-DAG: f32[1024]{0:S(7)}{{.*}}all-reduce{{.*}}channel_id=1 // CHECK-DAG: s32[1024]{0}{{.*}}all-reduce{{.*}}channel_id=2 // CHECK-DAG: f32[2048]{0}{{.*}}all-reduce{{.*}}channel_id=3 )"; @@ -3300,8 +3300,8 @@ TEST_F(GpuCompilerTest, SymmetricBuffersMultipleCollectives) { const HloModule* optimized_module = optimized_module_and_executable.first; constexpr absl::string_view expected_check = R"( - // CHECK-DAG: f32[1024]{0:S(1)}{{.*}}all-reduce{{.*}}channel_id=1 - // CHECK-DAG: f32[1024]{0:S(1)}{{.*}}all-gather{{.*}}channel_id=2 + // CHECK-DAG: f32[1024]{0:S(7)}{{.*}}all-reduce{{.*}}channel_id=1 + // CHECK-DAG: f32[1024]{0:S(7)}{{.*}}all-gather{{.*}}channel_id=2 // CHECK-DAG: f32[2048]{0}{{.*}}all-reduce{{.*}}channel_id=3 )"; @@ -3363,8 +3363,8 @@ TEST_F(GpuCompilerTest, SymmetricBuffersSeveralFilters) { const HloModule* optimized_module = optimized_module_and_executable.first; constexpr absl::string_view expected_check = R"( - // CHECK-DAG: f32[1024]{0:S(1)}{{.*}}all-reduce{{.*}}channel_id=1 - // CHECK-DAG: s32[2048]{0:S(1)}{{.*}}all-gather{{.*}}channel_id=2 + // CHECK-DAG: f32[1024]{0:S(7)}{{.*}}all-reduce{{.*}}channel_id=1 + // CHECK-DAG: s32[2048]{0:S(7)}{{.*}}all-gather{{.*}}channel_id=2 )"; EXPECT_THAT(RunFileCheck( @@ -3425,8 +3425,8 @@ TEST_F(GpuCompilerTest, SymmetricBuffersOverlappingFilters) { const HloModule* optimized_module = optimized_module_and_executable.first; constexpr absl::string_view expected_check = R"( - // CHECK-DAG: f32[1024]{0:S(1)}{{.*}}all-reduce{{.*}}channel_id=1 - // CHECK-DAG: f32[2048]{0:S(1)}{{.*}}all-reduce{{.*}}channel_id=2 + // CHECK-DAG: f32[1024]{0:S(7)}{{.*}}all-reduce{{.*}}channel_id=1 + // CHECK-DAG: f32[2048]{0:S(7)}{{.*}}all-reduce{{.*}}channel_id=2 )"; EXPECT_THAT(RunFileCheck( diff --git a/third_party/xla/xla/service/gpu/gpu_memory_space_assignment.cc b/third_party/xla/xla/service/gpu/gpu_memory_space_assignment.cc index c9eda097160eea..d598c075274480 100644 --- a/third_party/xla/xla/service/gpu/gpu_memory_space_assignment.cc +++ b/third_party/xla/xla/service/gpu/gpu_memory_space_assignment.cc @@ -69,6 +69,10 @@ const absl::NoDestructor> absl::StatusOr AsMemorySpaceColor(int64_t memory_space) { switch (memory_space) { + case 1: + // Legacy value for collective memory space before + // Layout::kCollectiveMemorySpace (7) was introduced + return MemorySpaceColor::kCollective; case static_cast(MemorySpaceColor::kDefault): case static_cast(MemorySpaceColor::kCollective): case static_cast(MemorySpaceColor::kTempBuffer): @@ -76,8 +80,9 @@ absl::StatusOr AsMemorySpaceColor(int64_t memory_space) { default: return InvalidArgument( "Invalid memory space %d. " - "Valid values are 0 (default), 1 (collective), 2 (temp).", - memory_space); + "Valid values are %d (default), %d (collective), %d (temp).", + memory_space, MemorySpaceColor::kDefault, + MemorySpaceColor::kCollective, MemorySpaceColor::kTempBuffer); } } @@ -265,8 +270,15 @@ absl::StatusOr DetermineBufferColor( // space from the layout. const HloPosition& defining_position = value->defining_position(); if (defining_position.shape().has_layout()) { - const BufferValue::Color memory_space = + BufferValue::Color memory_space = defining_position.shape().layout().memory_space(); + if (memory_space == 1) { + // Legacy value for collective memory space before + // Layout::kCollectiveMemorySpace (7) was introduced. + memory_space = + static_cast(MemorySpaceColor::kCollective); + } + if (memory_space != 0) { candidates.push_back(memory_space); } diff --git a/third_party/xla/xla/service/gpu/gpu_memory_space_assignment.h b/third_party/xla/xla/service/gpu/gpu_memory_space_assignment.h index 6fbda3c1cdd28a..8433f555458bca 100644 --- a/third_party/xla/xla/service/gpu/gpu_memory_space_assignment.h +++ b/third_party/xla/xla/service/gpu/gpu_memory_space_assignment.h @@ -23,6 +23,7 @@ limitations under the License. #include "absl/status/statusor.h" #include "absl/strings/string_view.h" #include "xla/hlo/ir/hlo_instruction.h" +#include "xla/layout.h" #include "xla/service/buffer_assignment.h" #include "xla/xla.pb.h" @@ -39,11 +40,11 @@ inline constexpr absl::string_view kResultsMemorySpacesAttr = enum class MemorySpaceColor { // Corresponds to stream_executor::MemoryTypes::kDefault or kUnified. // This memory can be allocated with any device allocation API. - kDefault = 0, + kDefault = Layout::kDefaultMemorySpace, // Corresponds to stream_executor::MemoryTypes::kCollective. This memory // should be compatible with symmetric memory requirements. - kCollective = 1, + kCollective = Layout::kCollectiveMemorySpace, // Temp buffers can be allocated within separate memory space (if // xla_gpu_temp_buffer_use_separate_color is set). This improves cuda-graphs diff --git a/third_party/xla/xla/service/gpu/gpu_memory_space_assignment_test.cc b/third_party/xla/xla/service/gpu/gpu_memory_space_assignment_test.cc index 814f38efdbd121..e2605cbe591816 100644 --- a/third_party/xla/xla/service/gpu/gpu_memory_space_assignment_test.cc +++ b/third_party/xla/xla/service/gpu/gpu_memory_space_assignment_test.cc @@ -441,7 +441,7 @@ TEST_F(GpuMemorySpaceAssignmentTest, HloModule m ENTRY main { - ROOT %custom-call.9 = (f16[8], f16[8]) custom-call(), custom_call_target="mosaic_gpu_v2", frontend_attributes={results_memory_spaces="{0:1}"} + ROOT %custom-call.9 = (f16[8], f16[8]) custom-call(), custom_call_target="mosaic_gpu_v2", frontend_attributes={results_memory_spaces="{0:7}"} } )"; @@ -474,14 +474,14 @@ TEST_F(GpuMemorySpaceAssignmentTest, } TEST(ParseIndexMemorySpacePairsTest, SinglePair) { - ASSERT_OK_AND_ASSIGN(auto pairs, ParseIndexMemorySpacePairs("{0:1}")); + ASSERT_OK_AND_ASSIGN(auto pairs, ParseIndexMemorySpacePairs("{0:7}")); ASSERT_EQ(pairs.size(), 1); EXPECT_EQ(pairs[0].first, 0); EXPECT_EQ(pairs[0].second, MemorySpaceColor::kCollective); } TEST(ParseIndexMemorySpacePairsTest, MultiplePairs) { - ASSERT_OK_AND_ASSIGN(auto pairs, ParseIndexMemorySpacePairs("{0:1,2:2}")); + ASSERT_OK_AND_ASSIGN(auto pairs, ParseIndexMemorySpacePairs("{0:7,2:2}")); ASSERT_EQ(pairs.size(), 2); EXPECT_EQ(pairs[0].first, 0); EXPECT_EQ(pairs[0].second, MemorySpaceColor::kCollective); @@ -496,7 +496,7 @@ TEST(ParseIndexMemorySpacePairsTest, EmptyBraces) { TEST(ParseIndexMemorySpacePairsTest, WhitespaceHandling) { ASSERT_OK_AND_ASSIGN(auto pairs, - ParseIndexMemorySpacePairs("{ 0 : 1 , 2 : 0 }")); + ParseIndexMemorySpacePairs("{ 0 : 7 , 2 : 0 }")); ASSERT_EQ(pairs.size(), 2); EXPECT_EQ(pairs[0].first, 0); EXPECT_EQ(pairs[0].second, MemorySpaceColor::kCollective); @@ -505,7 +505,7 @@ TEST(ParseIndexMemorySpacePairsTest, WhitespaceHandling) { } TEST(ParseIndexMemorySpacePairsTest, MissingBraces) { - EXPECT_FALSE(ParseIndexMemorySpacePairs("0:1").ok()); + EXPECT_FALSE(ParseIndexMemorySpacePairs("0:7").ok()); } TEST(ParseIndexMemorySpacePairsTest, InvalidPairFormat) { @@ -542,7 +542,7 @@ TEST_F(GpuMemorySpaceAssignmentTest, CustomCallOperandMemorySpace) { p1 = f32[1024]{0} parameter(1) ROOT custom-call = f32[1024]{0} custom-call(p0, p1), custom_call_target="my_custom_call", - frontend_attributes={operands_memory_spaces="{0:1}"} + frontend_attributes={operands_memory_spaces="{0:7}"} } )"; @@ -557,8 +557,9 @@ TEST_F(GpuMemorySpaceAssignmentTest, CustomCallOperandMemorySpace) { DependencyHloOrdering ordering(module.get()); EXPECT_OK(colorer(alias_analysis.get(), ordering)); - // Operand 0 (p0) should be colored with memory space 1. - EXPECT_EQ(FindColorByName(*alias_analysis, "p0"), 1); + // Operand 0 (p0) should be colored with memory space kCollective. + EXPECT_EQ(FindColorByName(*alias_analysis, "p0"), + (int)MemorySpaceColor::kCollective); // Operand 1 (p1) should remain in default memory space. EXPECT_EQ(FindColorByName(*alias_analysis, "p1"), (int)MemorySpaceColor::kDefault); @@ -572,7 +573,7 @@ TEST_F(GpuMemorySpaceAssignmentTest, CustomCallResultMemorySpace) { p0 = f32[1024]{0} parameter(0) ROOT custom-call = f32[1024]{0} custom-call(p0), custom_call_target="my_custom_call", - frontend_attributes={results_memory_spaces="{0:1}"} + frontend_attributes={results_memory_spaces="{0:7}"} } )"; @@ -587,8 +588,9 @@ TEST_F(GpuMemorySpaceAssignmentTest, CustomCallResultMemorySpace) { DependencyHloOrdering ordering(module.get()); EXPECT_OK(colorer(alias_analysis.get(), ordering)); - // The custom call result should be colored with memory space 1. - EXPECT_EQ(FindColorByName(*alias_analysis, "custom-call"), 1); + // The custom call result should be colored with memory space kCollective. + EXPECT_EQ(FindColorByName(*alias_analysis, "custom-call"), + (int)MemorySpaceColor::kCollective); } TEST_F(GpuMemorySpaceAssignmentTest, CustomCallTupleResultMemorySpace) { @@ -599,7 +601,7 @@ TEST_F(GpuMemorySpaceAssignmentTest, CustomCallTupleResultMemorySpace) { p0 = f32[1024]{0} parameter(0) custom-call = (f32[1024]{0}, f32[512]{0}) custom-call(p0), custom_call_target="my_custom_call", - frontend_attributes={results_memory_spaces="{0:1,1:1}"} + frontend_attributes={results_memory_spaces="{0:7,1:7}"} ROOT gte = f32[1024]{0} get-tuple-element(custom-call), index=0 } )"; @@ -621,8 +623,8 @@ TEST_F(GpuMemorySpaceAssignmentTest, CustomCallTupleResultMemorySpace) { if (value->instruction()->name() != "custom-call") continue; const ShapeIndex& idx = value->defining_index(); if (idx.size() == 1 && (idx[0] == 0 || idx[0] == 1)) { - EXPECT_EQ(value->color(), 1) - << "Tuple element " << idx[0] << " should have color 1"; + EXPECT_EQ(value->color(), (int)MemorySpaceColor::kCollective) + << "Tuple element " << idx[0] << " should have color kCollective"; } } } @@ -642,8 +644,8 @@ TEST_P(GpuMosaicCollectiveMemorySpaceAssignmentTest, HloModule m ENTRY main { ROOT %custom-call.9 = (f16[8])", - IsMosaicWithCollectiveMetadata() ? "{0:S(1)}" : "", R"(, f16[8])", - IsMosaicWithCollectiveMetadata() ? "{0:S(1)}" : "", + IsMosaicWithCollectiveMetadata() ? "{0:S(7)}" : "", R"(, f16[8])", + IsMosaicWithCollectiveMetadata() ? "{0:S(7)}" : "", R"() custom-call(), custom_call_target="mosaic_gpu_v2" } )"); diff --git a/third_party/xla/xla/service/hlo.proto b/third_party/xla/xla/service/hlo.proto index b3646dd5b84100..ea708ce2b8dc4c 100644 --- a/third_party/xla/xla/service/hlo.proto +++ b/third_party/xla/xla/service/hlo.proto @@ -128,7 +128,7 @@ enum CustomCallApiVersion { } // Serialization of HloInstruction. -// Next ID: 103 +// Next ID: 104 message HloInstructionProto { reserved 10; reserved "parameter_name"; @@ -442,6 +442,9 @@ message HloInstructionProto { // Comparison order (TOTAL, PARTIAL). string comparison_order = 102; + + // The mode of a kShuffle and the attributes of that mode. + xla.ShuffleMode shuffle_mode = 103; } // Serialization of HloComputation. diff --git a/third_party/xla/xla/service/hlo_cost_analysis.cc b/third_party/xla/xla/service/hlo_cost_analysis.cc index 11e310dfde8836..c646a5eae4a9a7 100644 --- a/third_party/xla/xla/service/hlo_cost_analysis.cc +++ b/third_party/xla/xla/service/hlo_cost_analysis.cc @@ -380,6 +380,10 @@ absl::Status HloCostAnalysis::HandleReverse(const HloInstruction*) { return absl::OkStatus(); } +absl::Status HloCostAnalysis::HandleShuffle(const HloInstruction*) { + return absl::OkStatus(); +} + absl::Status HloCostAnalysis::HandleSlice(const HloInstruction* slice) { const int64_t output_shape_size = GetShapeSize(slice->shape()); diff --git a/third_party/xla/xla/service/hlo_cost_analysis.h b/third_party/xla/xla/service/hlo_cost_analysis.h index cd97ea8832cc3a..ca809d4b6779ed 100644 --- a/third_party/xla/xla/service/hlo_cost_analysis.h +++ b/third_party/xla/xla/service/hlo_cost_analysis.h @@ -515,6 +515,7 @@ class HloCostAnalysis : public ConstDfsHloVisitor { absl::Status HandleRngGetAndUpdateState( const HloInstruction* random) override; absl::Status HandleReverse(const HloInstruction* reverse) override; + absl::Status HandleShuffle(const HloInstruction* shuffle) override; absl::Status HandleSort(const HloInstruction* sort) override; absl::Status HandleParameter(const HloInstruction* parameter) override; absl::Status HandleReduce(const HloInstruction* reduce) override; diff --git a/third_party/xla/xla/service/hlo_graph_dumper.cc b/third_party/xla/xla/service/hlo_graph_dumper.cc index ab7e51f833f368..55675bb8934fe6 100644 --- a/third_party/xla/xla/service/hlo_graph_dumper.cc +++ b/third_party/xla/xla/service/hlo_graph_dumper.cc @@ -1245,6 +1245,7 @@ ColorScheme HloDotDumper::GetInstructionColor(const HloInstruction* instr) { case HloOpcode::kReshape: case HloOpcode::kDynamicReshape: case HloOpcode::kReverse: + case HloOpcode::kShuffle: case HloOpcode::kTranspose: // These data-movement ops can be expensive; emphasize them. (Yes, even // concat can be expensive, at least on GPU, as it can create warp diff --git a/third_party/xla/xla/service/hlo_instruction_test.cc b/third_party/xla/xla/service/hlo_instruction_test.cc index 9d343dfc322867..8480299330cd5f 100644 --- a/third_party/xla/xla/service/hlo_instruction_test.cc +++ b/third_party/xla/xla/service/hlo_instruction_test.cc @@ -52,6 +52,7 @@ limitations under the License. #include "xla/service/pattern_matcher.h" #include "xla/shape.h" #include "xla/shape_util.h" +#include "xla/shuffle.h" #include "xla/tsl/util/proto/proto_matchers.h" #include "xla/util.h" #include "xla/window_util.h" @@ -1216,6 +1217,20 @@ TEST_F(HloInstructionTest, IdenticalInstructions) { Identical(*HloInstruction::CreateUnary(shape, HloOpcode::kCopy, op1), *HloInstruction::CreateUnary(shape, HloOpcode::kNegate, op1))); + // Shuffle. + EXPECT_TRUE(Identical(*HloInstruction::CreateShuffle(shape, op1, {0, 1}, + shuffle::Rotate({2, 3})), + *HloInstruction::CreateShuffle( + shape, op1, {0, 1}, shuffle::Rotate({2, 3})))); + EXPECT_FALSE(Identical(*HloInstruction::CreateShuffle( + shape, op1, {0, 1}, shuffle::Rotate({2, 3})), + *HloInstruction::CreateShuffle( + shape, op1, {0, 1}, shuffle::Rotate({3, 2})))); + EXPECT_FALSE(Identical(*HloInstruction::CreateShuffle( + shape, op1, {0, 1}, shuffle::Rotate({2, 3})), + *HloInstruction::CreateShuffle( + shape, op1, {1, 0}, shuffle::Rotate({2, 3})))); + // Tuples. EXPECT_TRUE(Identical(*HloInstruction::CreateTuple({op1, op2}), *HloInstruction::CreateTuple({op1, op2}))); diff --git a/third_party/xla/xla/service/hlo_verifier.cc b/third_party/xla/xla/service/hlo_verifier.cc index 431308d01822f7..0e532bfea16442 100644 --- a/third_party/xla/xla/service/hlo_verifier.cc +++ b/third_party/xla/xla/service/hlo_verifier.cc @@ -1190,6 +1190,14 @@ absl::Status ShapeVerifier::HandleReverse(HloInstruction* reverse) { reverse->dimensions())); } +absl::Status ShapeVerifier::HandleShuffle(HloInstruction* shuffle) { + HloShuffleInstruction* shuffle_instr = Cast(shuffle); + return CheckShape( + shuffle, ShapeInference::InferShuffleShape( + shuffle_instr->operand(0)->shape(), + shuffle_instr->dimensions(), shuffle_instr->shuffle_mode())); +} + absl::Status ShapeVerifier::HandleTopK(HloInstruction* hlo) { return CheckShape( hlo, ShapeInference::InferTopKShape(hlo->operand(0)->shape(), diff --git a/third_party/xla/xla/service/hlo_verifier.h b/third_party/xla/xla/service/hlo_verifier.h index 61f0d3e9ca8cf8..9af83c220ea9d3 100644 --- a/third_party/xla/xla/service/hlo_verifier.h +++ b/third_party/xla/xla/service/hlo_verifier.h @@ -253,6 +253,7 @@ class ShapeVerifier : public DfsHloVisitor { absl::Status HandleRngBitGenerator(HloInstruction*) override; absl::Status HandleRngGetAndUpdateState(HloInstruction*) override; absl::Status HandleReverse(HloInstruction* reverse) override; + absl::Status HandleShuffle(HloInstruction* shuffle) override; absl::Status HandleSort(HloInstruction* hlo) override; absl::Status HandleTopK(HloInstruction* hlo) override; absl::Status HandleConstant(HloInstruction* constant) override; diff --git a/third_party/xla/xla/service/hlo_verifier_test.cc b/third_party/xla/xla/service/hlo_verifier_test.cc index 60edc3cae66153..b459ed7ad29e8f 100644 --- a/third_party/xla/xla/service/hlo_verifier_test.cc +++ b/third_party/xla/xla/service/hlo_verifier_test.cc @@ -4403,6 +4403,76 @@ ENTRY TopK { ASSERT_TRUE(status.ok()); } +TEST_F(HloVerifierTest, ShuffleRotateOK) { + constexpr absl::string_view kHlo = R"( +HloModule shuffle, entry_computation_layout={(f32[10,20]{1,0})->f32[10,20]{1,0}} + +ENTRY Shuffle { + x = f32[10,20]{1,0} parameter(0) + ROOT shuffle = f32[10,20]{1,0} shuffle(x), dimensions={0,1}, mode=rotate, shifts={2,5} +} + )"; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnUnverifiedModule(kHlo)); + const absl::StatusOr result = verifier().Run(module.get()); + + EXPECT_THAT(result, absl_testing::IsOk()); +} + +TEST_F(HloVerifierTest, ShuffleInvalidDimensions) { + constexpr absl::string_view kHlo = R"( +HloModule shuffle, entry_computation_layout={(f32[10,20]{1,0})->f32[10,20]{1,0}} + +ENTRY Shuffle { + x = f32[10,20]{1,0} parameter(0) + ROOT shuffle = f32[10,20]{1,0} shuffle(x), dimensions={0,2}, mode=rotate, shifts={2,5} +} + )"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnUnverifiedModule(kHlo)); + const absl::StatusOr result = verifier().Run(module.get()); + + EXPECT_THAT(result, StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("out-of-bounds"))); +} + +TEST_F(HloVerifierTest, ShuffleRotateMismatchedShiftsSize) { + constexpr absl::string_view kHlo = R"( +HloModule shuffle, entry_computation_layout={(f32[10,20]{1,0})->f32[10,20]{1,0}} + +ENTRY Shuffle { + x = f32[10,20]{1,0} parameter(0) + ROOT shuffle = f32[10,20]{1,0} shuffle(x), dimensions={0,1}, mode=rotate, shifts={2} +} + )"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnUnverifiedModule(kHlo)); + const absl::StatusOr result = verifier().Run(module.get()); + + EXPECT_THAT(result, + StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("dimensions and shifts must have the same " + "size"))); +} + +TEST_F(HloVerifierTest, ShuffleDuplicatedDimensions) { + constexpr absl::string_view kHlo = R"( +HloModule shuffle, entry_computation_layout={(f32[10,20]{1,0})->f32[10,20]{1,0}} + +ENTRY Shuffle { + x = f32[10,20]{1,0} parameter(0) + ROOT shuffle = f32[10,20]{1,0} shuffle(x), dimensions={0,0}, mode=rotate, shifts={2,5} +} + )"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnUnverifiedModule(kHlo)); + const absl::StatusOr result = verifier().Run(module.get()); + + EXPECT_THAT(result, StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("duplicated"))); +} + TEST_F(HloVerifierTest, InputLayoutMismatchIgnored) { // Note: The mismatch is between the entry_computation_layout and the layout // of parameter(1). diff --git a/third_party/xla/xla/service/instruction_fusion.cc b/third_party/xla/xla/service/instruction_fusion.cc index 7c900b9629562c..b32a0ee94944f0 100644 --- a/third_party/xla/xla/service/instruction_fusion.cc +++ b/third_party/xla/xla/service/instruction_fusion.cc @@ -132,6 +132,7 @@ bool IsAlwaysDuplicable(const HloInstruction& instruction) { case HloOpcode::kShiftLeft: case HloOpcode::kShiftRightArithmetic: case HloOpcode::kShiftRightLogical: + case HloOpcode::kShuffle: case HloOpcode::kSlice: case HloOpcode::kStochasticConvert: case HloOpcode::kSubtract: diff --git a/third_party/xla/xla/service/layout_assignment.cc b/third_party/xla/xla/service/layout_assignment.cc index 425e2f46e154bd..88fa67ba987317 100644 --- a/third_party/xla/xla/service/layout_assignment.cc +++ b/third_party/xla/xla/service/layout_assignment.cc @@ -3820,6 +3820,7 @@ bool LayoutAssignment::InstructionCanChangeLayout( case HloOpcode::kShiftLeft: case HloOpcode::kShiftRightArithmetic: case HloOpcode::kShiftRightLogical: + case HloOpcode::kShuffle: case HloOpcode::kSign: case HloOpcode::kSin: case HloOpcode::kSinh: diff --git a/third_party/xla/xla/service/memory_space_assignment/allocation.cc b/third_party/xla/xla/service/memory_space_assignment/allocation.cc new file mode 100644 index 00000000000000..7e060854083011 --- /dev/null +++ b/third_party/xla/xla/service/memory_space_assignment/allocation.cc @@ -0,0 +1,1502 @@ +/* Copyright 2024 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xla/service/memory_space_assignment/allocation.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "absl/algorithm/container.h" +#include "absl/base/casts.h" +#include "absl/container/flat_hash_map.h" +#include "absl/container/flat_hash_set.h" +#include "absl/functional/function_ref.h" +#include "absl/log/check.h" +#include "absl/log/log.h" +#include "absl/status/status.h" +#include "absl/status/status_macros.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/str_join.h" +#include "absl/strings/string_view.h" +#include "absl/types/span.h" +#include "xla/hlo/analysis/hlo_alias_analysis.h" +#include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/ir/hlo_original_value.h" +#include "xla/hlo/ir/hlo_original_value_util.h" +#include "xla/hlo/utils/hlo_live_range.h" +#include "xla/layout.h" +#include "xla/layout_util.h" +#include "xla/service/heap_simulator/allocation_block.h" +#include "xla/service/heap_simulator/heap_simulator.h" +#include "xla/service/hlo_value.h" +#include "xla/service/memory_space_assignment/memory_space_assignment.pb.h" +#include "xla/service/memory_space_assignment/slice.h" +#include "xla/service/time_utils.h" +#include "xla/service/tuple_util.h" +#include "xla/service/while_util.h" +#include "xla/shape.h" +#include "xla/shape_util.h" +#include "xla/status_macros.h" +#include "xla/tsl/platform/errors.h" +#include "xla/tsl/platform/statusor.h" +#include "xla/util.h" +#include "xla/xla_data.pb.h" +#include "tsl/platform/casts.h" +#include "tsl/platform/errors.h" +#include "tsl/platform/statusor.h" + +namespace xla::memory_space_assignment { +namespace { + +std::string UsesToString(const std::vector& uses) { + if (uses.empty()) { + return "none"; + } + std::vector uses_str; + uses_str.reserve(uses.size()); + for (const auto& use : uses) { + uses_str.push_back(use.ToString()); + } + return absl::StrJoin(uses_str, ","); +} + +void AppendToTupleOriginalValue(const HloInstruction* src_tuple, + HloInstruction* dest_tuple, + const HloInstruction* appended_instr) { + if (src_tuple == nullptr || dest_tuple == nullptr || + !dest_tuple->shape().IsTuple() || + dest_tuple->shape().tuple_shapes().empty()) { + return; + } + bool has_appended_ov = + appended_instr != nullptr && appended_instr->original_value() != nullptr; + if (src_tuple->original_value() == nullptr && !has_appended_ov) { + return; + } + int64_t new_tuple_index = + static_cast(dest_tuple->shape().tuple_shapes().size()) - 1; + if (src_tuple->original_value() != nullptr) { + int64_t src_size = src_tuple->original_value()->tree().num_children(); + int64_t num_to_copy = std::min(src_size, new_tuple_index); + absl::flat_hash_map old_to_new_tuple_idx; + for (int64_t i = 0; i < num_to_copy; ++i) { + old_to_new_tuple_idx[i] = i; + } + CopyOriginalValue(src_tuple, dest_tuple, old_to_new_tuple_idx); + } else { + dest_tuple->set_original_value( + std::make_shared(dest_tuple->shape())); + } + if (has_appended_ov && dest_tuple->original_value() != nullptr) { + CHECK_OK( + dest_tuple->original_value()->mutable_tree()->CopyCompatibleSubtreeFrom( + appended_instr->original_value()->tree(), {}, {new_tuple_index})) + << "Incompatible OriginalValue subtree for appended element at index " + << new_tuple_index << ": " << appended_instr->ToString(); + } + CHECK(dest_tuple->original_value()->IsCompatibleWith(dest_tuple->shape())) + << "Updated OriginalValue is incompatible with shape: " + << dest_tuple->ToString(); +} + +// Helper function to compute the start time for a SlicedCopyAllocation. +int64_t GetSlicedCopyAllocationExclusiveStartTime( + const std::vector& + slice_decisions_sorted_by_exclusive_start_time) { + if (slice_decisions_sorted_by_exclusive_start_time.empty()) { + return -1; + } + + return slice_decisions_sorted_by_exclusive_start_time.front() + .exclusive_start_time; +} + +// Helper function to compute the underlying Allocation chunk for a +// SlicedCopyAllocation. +std::optional GetSlicedCopyAllocationChunk( + const std::vector& slice_decisions_sorted_by_start_time) { + if (slice_decisions_sorted_by_start_time.empty()) { + return std::nullopt; + } + auto offset_cmp = [](const SliceDecision& lhs, const SliceDecision& rhs) { + return lhs.chunk.offset < rhs.chunk.offset; + }; + auto end_cmp = [](const SliceDecision& lhs, const SliceDecision& rhs) { + return lhs.chunk.chunk_end() < rhs.chunk.chunk_end(); + }; + return HeapSimulator::Chunk::FromOffsetEnd( + std::min_element(slice_decisions_sorted_by_start_time.begin(), + slice_decisions_sorted_by_start_time.end(), offset_cmp) + ->chunk.offset, + std::max_element(slice_decisions_sorted_by_start_time.begin(), + slice_decisions_sorted_by_start_time.end(), end_cmp) + ->chunk.chunk_end()); +} + +} // namespace + +std::string MemorySpaceToString(MemorySpace memory_space) { + switch (memory_space) { + case MemorySpace::kDefault: + return "def"; + case MemorySpace::kAlternate: + return "alt"; + } +} + +std::optional Allocation::cross_program_prefetch_index() const { + return cross_program_prefetch_index_; +} + +HeapSimulator::Chunk Allocation::chunk() const { + CHECK(chunk_.has_value()); + return *chunk_; +} + +void Allocation::set_offset(int64_t offset) { + CHECK(chunk_.has_value()); + *chunk_ = HeapSimulator::Chunk::FromOffsetSize(offset, chunk_->size); +} + +bool Allocation::is_in_alternate_mem() const { + return memory_space_ == MemorySpace::kAlternate; +} + +bool Allocation::is_in_default_mem() const { + return memory_space_ == MemorySpace::kDefault; +} + +void Allocation::RemoveUse(HloUse use) { + uses_.erase(std::remove_if(uses_.begin(), uses_.end(), + [=](const auto& u) { return u == use; }), + uses_.end()); +} + +void Allocation::AddUse(HloUse use) { + HloInstruction* operand = + use.instruction->mutable_operand(use.operand_number); + // If the use is a tuple, look inside the tuple to find the actual use. + for (int64_t index : use.operand_index) { + if (operand->opcode() != HloOpcode::kTuple) { + break; + } + operand = operand->mutable_operand(index); + } + + // Look beyond GetTupleElement(Tuple()) pattern for any bitcasts. + std::function get_simplified_operand; + get_simplified_operand = [&](HloInstruction* instruction) { + while (instruction->opcode() == HloOpcode::kGetTupleElement) { + HloInstruction* operand = + get_simplified_operand(instruction->mutable_operand(0)); + if (operand->opcode() == HloOpcode::kTuple) { + instruction = operand->mutable_operand(instruction->tuple_index()); + } else { + return instruction; + } + } + return instruction; + }; + operand = get_simplified_operand(operand); + + uses_.push_back(use); +} + +absl::Status Allocation::UpdateUses(HloComputation* computation, + HloInstruction* producing_instruction, + const BitcastSplitFn& bitcast_split_fn, + const HloLiveRange& hlo_live_range, + const HloAliasAnalysis& alias_analysis) { + for (const HloUse& use : uses()) { + HloInstruction* replacement_instruction = producing_instruction; + const Shape& operand_shape = + use.instruction->operand(use.operand_number)->shape(); + if (operand_shape.IsTuple()) { + HloInstruction* tuple_inst = + use.instruction->mutable_operand(use.operand_number); + + // Trace back to original instructions in the schedule to use + // HloAliasAnalysis safely. + auto find_original_leaf = [&](const HloInstruction* inst, + ShapeIndex index) + -> std::pair { + while (!hlo_live_range.instruction_schedule().contains(inst)) { + if (inst->opcode() == HloOpcode::kTuple) { + if (index.empty()) { + break; + } + int64_t op_idx = index[0]; + inst = inst->operand(op_idx); + index = ShapeIndex(index.begin() + 1, index.end()); + } else if (inst->opcode() == HloOpcode::kGetTupleElement) { + ShapeIndex new_index; + new_index.push_back(inst->tuple_index()); + for (int i : index) { + new_index.push_back(i); + } + inst = inst->operand(0); + index = new_index; + } else if (inst->opcode() == HloOpcode::kBitcast || + inst->opcode() == HloOpcode::kCopy || + inst->opcode() == HloOpcode::kCopyDone || + inst->opcode() == HloOpcode::kCopyStart) { + inst = inst->operand(0); + } else { + break; + } + } + return {inst, index}; + }; + + bool should_skip_reconstruction = false; + + // If the producing instruction shares the same buffer with the operand, + // then we skip the reconstruction. + auto is_aliasing_operand = [&](int operand_index) { + ShapeIndex sub_index; + if (use.operand_index.size() > 1) { + sub_index = ShapeIndex(use.operand_index.begin() + 1, + use.operand_index.end()); + } + auto [traced_operand, traced_index] = + find_original_leaf(tuple_inst->operand(operand_index), sub_index); + return hlo_live_range.instruction_schedule().contains(traced_operand) && + hlo_live_range.instruction_schedule().contains(tuple_inst) && + alias_analysis.GetUniqueBufferAt(traced_operand, traced_index) + .id() == + alias_analysis + .GetUniqueBufferAt(tuple_inst, use.operand_index) + .id(); + }; + + for (int operand_index = 0; operand_index < tuple_inst->operand_count(); + ++operand_index) { + if (tuple_inst->opcode() == HloOpcode::kTuple && + !use.operand_index.empty() && + operand_index != use.operand_index.front()) { + continue; + } + if (tuple_inst->operand(operand_index) == producing_instruction && + is_aliasing_operand(operand_index)) { + should_skip_reconstruction = true; + break; + } + } + + if (should_skip_reconstruction) { + VLOG(4) << "Skipping tuple reconstruction for " << tuple_inst->name() + << " at index " << use.operand_index.ToString(); + replacement_instruction = tuple_inst; + } else { + ABSL_ASSIGN_OR_RETURN( + replacement_instruction, + TupleUtil::ReplaceTupleWith(producing_instruction, tuple_inst, + use.operand_index)); + if (replacement_instruction != tuple_inst) { + replacement_instruction->CopyOriginalValue(tuple_inst, + /*clone=*/true); + if (replacement_instruction->original_value() != nullptr && + producing_instruction->original_value() != nullptr) { + CHECK_OK(replacement_instruction->original_value() + ->mutable_tree() + ->CopyCompatibleSubtreeFrom( + producing_instruction->original_value()->tree(), + {}, use.operand_index)); + } + } + } + } else if (!Shape::Equal().IgnoreSplitConfigInLayout()( + operand_shape, producing_instruction->shape())) { + // When processing allocations, we treat bitcasts as trivial positions and + // do not create allocations for them. We insert bitcasts after copies, to + // account for the fact that we don't have an allocation for the bitcast. + VLOG(4) << "Old shape = " << operand_shape.ToString() + << ", new shape = " << producing_instruction->shape().ToString() + << "; inserting a bitcast."; + replacement_instruction = computation->AddInstruction( + HloInstruction::CreateBitcast(operand_shape, producing_instruction)); + replacement_instruction->CopyOriginalValue(producing_instruction); + if (mutable_split_shape().has_value() && + producing_instruction->shape().layout().split_configs().size() != 0) { + ABSL_ASSIGN_OR_RETURN( + int64_t split_dim, + bitcast_split_fn( + replacement_instruction, + mutable_split_shape()->layout().split_configs(0).dimension())); + SplitConfig split_config( + split_dim, + {mutable_split_shape()->layout().split_configs(0).split_indices( + 0)}); + replacement_instruction->mutable_shape() + ->mutable_layout() + ->add_split_configs(split_config); + } + } + ABSL_RETURN_IF_ERROR(use.instruction->ReplaceOperandWith( + use.operand_number, replacement_instruction)); + } + return absl::OkStatus(); +} + +bool Allocation::is_copy_like_allocation() const { + return is_copy_allocation() || is_sliced_copy_allocation(); +} + +HloInstruction* Allocation::AddGetTupleElements() const { + CHECK_NE(defining_position().instruction, nullptr); + + Shape shape = defining_position().shape(); + CHECK(shape.IsArray()) << "Allocation shape is not an array. Shape = " + << shape.ToString() + << " position = " << defining_position().shape(); + HloInstruction* result = TupleUtil::AddGetTupleElements(defining_position()); + if (result != defining_position().instruction && + result->original_value() == nullptr && + defining_position().instruction->original_value() != nullptr) { + std::shared_ptr src_ov = + defining_position().instruction->original_value(); + const auto& src_tree = src_ov->tree(); + if (src_tree.find(defining_position().index) != src_tree.end()) { + auto new_ov = std::make_shared(result->shape()); + CHECK_OK(new_ov->mutable_tree()->CopyCompatibleSubtreeFrom( + src_tree, defining_position().index, {})); + if (!new_ov->IsEmpty()) { + result->set_original_value(std::move(new_ov)); + } + } + } + return result; +} + +Allocation::Allocation(HloPosition defining_position, MemorySpace memory_space, + std::optional chunk, + int64_t start_time, int64_t end_time, + std::optional cross_program_prefetch_index) + : original_defining_position_(std::move(defining_position)), + memory_space_(memory_space), + chunk_(chunk), + start_time_(start_time), + end_time_(end_time), + cross_program_prefetch_index_(cross_program_prefetch_index), + split_shape_(std::nullopt) {} + +HloPosition Allocation::original_defining_position() const { + return original_defining_position_; +} + +void Allocation::set_original_defining_position(HloPosition defining_position) { + original_defining_position_ = std::move(defining_position); +} + +bool Allocation::base_is_equal(const Allocation& other) const { + return defining_position() == other.defining_position() && + uses() == other.uses() && memory_space() == other.memory_space() && + chunk() == other.chunk() && start_time() == other.start_time() && + end_time() == other.end_time() && + earliest_available_time() == other.earliest_available_time(); +} + +PinnedAllocation::PinnedAllocation(HloPosition defining_position, + MemorySpace memory_space, + std::optional chunk, + int64_t start_time, int64_t end_time) + : Allocation(std::move(defining_position), memory_space, chunk, start_time, + end_time, + /*cross_program_prefetch_index=*/std::nullopt) {} + +HloPosition PinnedAllocation::defining_position() const { + return original_defining_position(); +} + +bool PinnedAllocation::operator==(const PinnedAllocation& other) const { + return this->base_is_equal(static_cast(other)); +} + +bool MirroredAllocation::operator==(const MirroredAllocation& other) const { + return this->base_is_equal(static_cast(other)); +} + +bool ParentAllocation::operator==(const ParentAllocation& other) const { + return this->base_is_equal(static_cast(other)); +} + +bool PinnedAllocation::operator==(const Allocation& other) const { + const PinnedAllocation* casted_other = + dynamic_cast(&other); + return casted_other != nullptr && (*this) == (*casted_other); +} + +absl::Status PinnedAllocation::Process(const BitcastSplitFn& bitcast_split_fn, + const HloLiveRange& hlo_live_range, + const HloAliasAnalysis& alias_analysis) { + HloInstruction* producing_instruction = AddGetTupleElements(); + HloComputation* computation = producing_instruction->parent(); + return UpdateUses(computation, producing_instruction, bitcast_split_fn, + hlo_live_range, alias_analysis); +} + +std::string PinnedAllocation::ToString() const { + std::string memory_space_str = MemorySpaceToString(memory_space()); + std::optional chunk = maybe_chunk(); + if (chunk) { + absl::StrAppend(&memory_space_str, " (off: ", chunk->offset, + ", size: ", chunk->size, ")"); + } + return absl::StrCat("PinnedAllocation in ", memory_space_str, " defined at ", + original_defining_position().ToString(), + ", start_time:", start_time(), ", end_time:", end_time(), + ", uses: ", UsesToString(uses())); +} + +void PinnedAllocation::MarkIfNeeded( + absl::flat_hash_set& needed_allocations) const { + MarkNeeded(needed_allocations); +} + +void PinnedAllocation::MarkNeeded( + absl::flat_hash_set& needed_allocations) const { + needed_allocations.insert(this); +} + +ReservedAllocation::ReservedAllocation(HloPosition defining_position, + HeapSimulator::Chunk chunk, + int64_t start_time, int64_t end_time) + : Allocation(std::move(defining_position), MemorySpace::kAlternate, chunk, + /*start_time=*/start_time, /*end_time=*/end_time, + /*cross_program_prefetch_index=*/std::nullopt), + reserved_(true) {} + +HloPosition ReservedAllocation::defining_position() const { + return original_defining_position(); +} + +absl::Status ReservedAllocation::Process( + const BitcastSplitFn& bitcast_split_fn, const HloLiveRange& hlo_live_range, + const HloAliasAnalysis& alias_analysis) { + return absl::OkStatus(); +} + +void ReservedAllocation::MarkIfNeeded( + absl::flat_hash_set& needed_allocations) const { + MarkNeeded(needed_allocations); +} + +void ReservedAllocation::MarkNeeded( + absl::flat_hash_set& needed_allocations) const { + needed_allocations.insert(this); +} + +std::string ReservedAllocation::ToString() const { + std::string memory_space_str = MemorySpaceToString(memory_space()); + std::optional chunk = maybe_chunk(); + if (chunk) { + absl::StrAppend(&memory_space_str, " chunk: ", chunk->ToString()); + } + return absl::StrCat("ReservedAllocation defined in ", memory_space_str, + " defined at ", original_defining_position().ToString(), + ", start_time: ", start_time(), + ", end_time: ", end_time(), " reserved: ", reserved_); +} + +bool ReservedAllocation::operator==(const Allocation& other) const { + const ReservedAllocation* casted_other = + dynamic_cast(&other); + return casted_other != nullptr && (*this) == (*casted_other); +} + +bool ReservedAllocation::operator==(const ReservedAllocation& other) const { + return this->base_is_equal(static_cast(other)) && + reserved_ == other.reserved_; +} + +CopyAllocation::CopyAllocation( + Allocation& prev_allocation, MemorySpace memory_space, + std::optional chunk, + int64_t copy_start_schedule_after_time, + int64_t copy_done_schedule_before_time, int64_t end_time, + std::optional cross_program_prefetch_index, + HloInstruction* sync_mem_op, HloInstruction* async_mem_op_start, + HloInstruction* async_mem_op_done, int64_t source_operand_index) + : Allocation( + /*defining_position=*/{nullptr, {}}, memory_space, chunk, + // Allocation uses an inclusive start time + ExclusiveToInclusiveStartTime(copy_start_schedule_after_time), + end_time, cross_program_prefetch_index), + prev_allocation_(prev_allocation), + copy_start_schedule_after_(copy_start_schedule_after_time), + copy_done_schedule_before_(copy_done_schedule_before_time), + copy_start_(async_mem_op_start), + copy_done_(async_mem_op_done), + sync_mem_op_(sync_mem_op), + source_operand_index_(source_operand_index) {} + +int64_t CopyAllocation::earliest_available_time() const { + return copy_done_schedule_before_; +} + +absl::Status CopyAllocation::Process(const BitcastSplitFn& bitcast_split_fn, + const HloLiveRange& hlo_live_range, + const HloAliasAnalysis& alias_analysis) { + // If the copy start and copy done instructions are already present in the + // original schedule we just need to schedule them in correct positions. There + // is no need to create new instructions. + if (copy_start_ != nullptr && copy_done_ != nullptr) { + // Update the allocation position with the copy complete instruction, so + // that if there are further copies from it, they can find the correct + // position. + set_original_defining_position(HloPosition{copy_start_, {0}}); + return absl::OkStatus(); + } + + // Copy allocations need to insert asynchronous copy nodes. + Shape shape = defining_position().shape(); + HloInstruction* producing_instruction = AddGetTupleElements(); + HloComputation* computation = producing_instruction->parent(); + VLOG(3) << "Processing copy allocation: " << ToString(); + if (sync_mem_op_ != nullptr && sync_mem_op_->opcode() != HloOpcode::kCopy) { + if (sync_mem_op_->opcode() == HloOpcode::kSlice || + sync_mem_op_->opcode() == HloOpcode::kDynamicSlice || + sync_mem_op_->IsCustomFusion()) { + ABSL_ASSIGN_OR_RETURN(copy_done_, + computation->CreateAsyncInstructions( + sync_mem_op_, {ShapeUtil::MakeShape(S32, {})}, + HloInstruction::kMainExecutionThread, false)); + } else { + return Internal("Sync mem op is not a copy, slice, or dynamic slice."); + } + copy_start_ = copy_done_->mutable_operand(0); + // If the shape of the copy start operand is not compatible with the + // shape of the producing instruction, we insert a bitcast to make them + // compatible. + if (!ShapeUtil::CompatibleIgnoringFpPrecision( + producing_instruction->shape(), + copy_start_->operand(source_operand_index_)->shape())) { + HloInstruction* bitcast = + computation->AddInstruction(HloInstruction::CreateBitcast( + copy_start_->operand(source_operand_index_)->shape(), + producing_instruction)); + bitcast->CopyOriginalValue(producing_instruction); + producing_instruction = bitcast; + } + ABSL_RETURN_IF_ERROR(copy_start_->ReplaceOperandWith(source_operand_index_, + producing_instruction)); + } else { + Shape dest_shape = shape; + if (memory_space() == MemorySpace::kDefault) { + dest_shape.mutable_layout()->clear_split_configs(); + } + copy_start_ = computation->AddInstruction(HloInstruction::CreateCopyStart( + ShapeUtil::MakeTupleShape( + {dest_shape, shape, ShapeUtil::MakeShape(U32, {})}), + producing_instruction, cross_program_prefetch_index())); + copy_done_ = computation->AddInstruction(HloInstruction::CreateUnary( + dest_shape, HloOpcode::kCopyDone, copy_start_)); + } + copy_done_->CopyOriginalValue(producing_instruction); + VLOG(4) << "Created " << copy_start_->name() + << " for copy allocation: " << ToString(); + + // Update the allocation position with the copy complete instruction, so that + // if there are further copies from it, they can find the correct position. + set_original_defining_position(HloPosition{copy_done_, {}}); + return UpdateUses(computation, copy_done_, bitcast_split_fn, hlo_live_range, + alias_analysis); +} + +void CopyAllocation::MarkIfNeeded( + absl::flat_hash_set& needed_allocations) const { + MarkNeeded(needed_allocations); +} + +void CopyAllocation::MarkNeeded( + absl::flat_hash_set& needed_allocations) const { + if (needed_allocations.insert(this).second) { + prev_allocation_.MarkNeeded(needed_allocations); + } +} + +std::string CopyAllocation::ToString() const { + std::string memory_space_str = MemorySpaceToString(memory_space()); + std::optional chunk = maybe_chunk(); + if (chunk) { + absl::StrAppend(&memory_space_str, " (off: ", chunk->offset, + ", size: ", chunk->size, ")"); + } + return absl::StrCat( + "Copy Allocation in ", memory_space_str, ", start_time:", start_time(), + ", end_time:", end_time(), + ", copy_start_after_time: ", copy_start_schedule_after(), + ", copy_done_before_time: ", copy_done_schedule_before(), + ", uses: ", UsesToString(uses()), + ", existing sync_mem_op: ", sync_mem_op_ ? sync_mem_op_->name() : "none", + ", existing copy_start: ", copy_start_ ? copy_start_->name() : "none", + ", existing copy_done: ", copy_done_ ? copy_done_->name() : "none", + ", from ", prev_allocation_.ToString()); +} + +HloPosition CopyAllocation::defining_position() const { + // Unless explicitly set, the defining position of a copy allocation is + // retrieved from the previous allocation. This is because we don't create + // new CopyStart/CopyDone instructions until later and the position should + // point to the previous (copy or otherwise) allocation's position for the + // original defining position. + HloPosition defining_position = original_defining_position(); + if (defining_position.instruction == nullptr) { + return prev_allocation_.defining_position(); + } + return defining_position; +} + +bool CopyAllocation::operator==(const CopyAllocation& other) const { + return this->base_is_equal(static_cast(other)) && + copy_done_schedule_before() == other.copy_done_schedule_before() && + copy_start_schedule_after() == other.copy_start_schedule_after() && + copy_start() == other.copy_start() && copy_done() == other.copy_done(); +} + +bool CopyAllocation::operator==(const Allocation& other) const { + const CopyAllocation* casted_other = + dynamic_cast(&other); + return casted_other != nullptr && (*this) == (*casted_other); +} + +void CopyAllocation::set_copy_start_schedule_after( + int64_t copy_start_schedule_after) { + copy_start_schedule_after_ = copy_start_schedule_after; +} + +void CopyAllocation::set_copy_done_schedule_before( + int64_t copy_done_schedule_before) { + copy_done_schedule_before_ = copy_done_schedule_before; +} + +int64_t CopyAllocation::copy_start_schedule_after() const { + return copy_start_schedule_after_; +} + +int64_t CopyAllocation::copy_done_schedule_before() const { + return copy_done_schedule_before_; +} + +SlicedCopyAllocation::SlicedCopyAllocation( + const Allocation& prev_allocation, MemorySpace memory_space, + std::vector slice_decisions_sorted_by_exclusive_start_time, + int64_t copy_done_schedule_before_time, int64_t end_time, + const SlicedPrefetchOptions& sliced_prefetch_options, + absl::FunctionRef get_equivalent_s8_shape_fn, + HloInstruction* sync_mem_op) + : Allocation( + /*defining_position=*/{nullptr, {}}, memory_space, + GetSlicedCopyAllocationChunk( + slice_decisions_sorted_by_exclusive_start_time), + // Allocation uses an inclusive start time + ExclusiveToInclusiveStartTime( + GetSlicedCopyAllocationExclusiveStartTime( + slice_decisions_sorted_by_exclusive_start_time)), + end_time, + /*cross_program_prefetch_index=*/std::nullopt), + original_shape_to_slice_(prev_allocation.defining_position().shape()), + prev_allocation_(prev_allocation), + sliced_prefetch_options_(sliced_prefetch_options), + get_equivalent_s8_shape_fn_(get_equivalent_s8_shape_fn), + sync_mem_op_(sync_mem_op) { + CHECK_GE(slice_decisions_sorted_by_exclusive_start_time.size(), 2); + slice_details_sorted_by_exclusive_start_time_.reserve( + slice_decisions_sorted_by_exclusive_start_time.size()); + for (SliceDecision& decision : + slice_decisions_sorted_by_exclusive_start_time) { + int64_t copy_done_schedule_after_time = decision.exclusive_start_time; + slice_details_sorted_by_exclusive_start_time_.push_back(SliceDetail{ + std::move(decision), + copy_done_schedule_after_time, + copy_done_schedule_before_time, + /*copy_start=*/nullptr, + /*copy_done=*/nullptr, + }); + } +} + +absl::Status SlicedCopyAllocation::Process( + const BitcastSplitFn& bitcast_split_fn, const HloLiveRange& hlo_live_range, + const HloAliasAnalysis& alias_analysis) { + Shape shape = defining_position().shape(); + HloInstruction* producing_instruction = AddGetTupleElements(); + + // Calling Process() over the previous allocation might + // have modified the defining position, and hence the shape that was used when + // we computed the slices. In cases where the shape has changed, we insert a + // bitcast, so slice instructions operate on the originally sliced shape. + // + // Note, these bitcasts are being inserted in the same cases that + // UpdateUses() is inserting bitcasts, except we are + // inserting the bitcasts before the copy, instead of after the copy. + if (!Shape::Equal().IgnoreMemorySpaceInLayout()(shape, + original_shape_to_slice_)) { + int64_t new_memory_space = shape.layout().memory_space(); + shape = original_shape_to_slice_; + shape.mutable_layout()->set_memory_space(new_memory_space); + producing_instruction = producing_instruction->parent()->AddInstruction( + HloInstruction::CreateBitcast(shape, producing_instruction)); + } + + HloComputation* computation = producing_instruction->parent(); + std::vector slice_dones; + slice_dones.reserve(slice_details_sorted_by_exclusive_start_time_.size()); + + // If we are trying to make all slices a uniform size, we bitcast the + // producing instruction to an array of bytes, so it is easy to slice into any + // size. + Shape slice_shape = shape; + if (IsUniformSliceSizingEnabled(sliced_prefetch_options_)) { + slice_shape = get_equivalent_s8_shape_fn_(shape); + producing_instruction = producing_instruction->parent()->AddInstruction( + HloInstruction::CreateBitcast(slice_shape, producing_instruction)); + } + + // Sliced copy allocations need to insert asynchronous copy nodes. + for (SliceDetail& slice_detail : + slice_details_sorted_by_exclusive_start_time_) { + ABSL_RETURN_IF_ERROR(slice_detail.CreateAsyncSlice( + slice_shape, *producing_instruction, *computation)); + VLOG(4) << "Created " << slice_detail.copy_start->name() + << " for sliced copy allocation: " << ToString(); + slice_dones.push_back(slice_detail.copy_done); + } + + ABSL_RETURN_IF_ERROR(CreateBitcastConcat(shape, slice_dones)); + + // If we bitcast to an array of bytes above, the result of the concatenated + // slices will also be an array of bytes. Thus, we need to cast the + // concatenation back to the original shape. + if (IsUniformSliceSizingEnabled(sliced_prefetch_options_)) { + concat_ = concat_->parent()->AddInstruction( + HloInstruction::CreateBitcast(shape, concat_)); + } + + // Update the allocation position with the copy complete instruction, so that + // if there are further copies from it, they can find the correct position. + set_original_defining_position(HloPosition{concat_, {}}); + return UpdateUses(computation, concat_, bitcast_split_fn, hlo_live_range, + alias_analysis); +} + +void SlicedCopyAllocation::MarkIfNeeded( + absl::flat_hash_set& needed_allocations) const { + MarkNeeded(needed_allocations); +} + +void SlicedCopyAllocation::MarkNeeded( + absl::flat_hash_set& needed_allocations) const { + if (needed_allocations.insert(this).second) { + prev_allocation_.MarkNeeded(needed_allocations); + } +} + +HloPosition SlicedCopyAllocation::defining_position() const { + // Unless explicitly set, the defining position of a sliced copy allocation is + // retrieved from the previous allocation. This is because we don't create + // new CopyStart/CopyDone instructions until later and the position should + // point to the previous (copy or otherwise) allocation's position for the + // original defining position. + HloPosition defining_position = original_defining_position(); + if (defining_position.instruction == nullptr) { + return prev_allocation_.defining_position(); + } + return defining_position; +} + +int64_t SlicedCopyAllocation::earliest_available_time() const { + return slice_details_sorted_by_start_time().back().copy_done_before_time; +} + +std::vector SlicedCopyAllocation::SliceOffsetsSortedByStartTime() + const { + std::vector offsets; + offsets.reserve(slice_details_sorted_by_exclusive_start_time_.size()); + + for (const SliceDetail& slice_detail : + slice_details_sorted_by_exclusive_start_time_) { + offsets.push_back(slice_detail.slice_decision.chunk.offset); + } + + return offsets; +} + +void SlicedCopyAllocation::AddDiffToAllSliceOffsets(int64_t diff) { + for (SliceDetail& slice_detail : + slice_details_sorted_by_exclusive_start_time_) { + HeapSimulator::Chunk& chunk = slice_detail.slice_decision.chunk; + chunk = + HeapSimulator::Chunk::FromOffsetSize(chunk.offset + diff, chunk.size); + } +} + +void SlicedCopyAllocation::ImportRepackedSliceData( + const SlicedAllocationData& data) { + int num_slices = slice_details_sorted_by_exclusive_start_time_.size(); + CHECK_EQ(data.slices_sorted_by_offset.size(), num_slices); + + std::vector slice_details_sorted_by_offset; + slice_details_sorted_by_offset.reserve(num_slices); + for (SliceDetail& slice_detail : + slice_details_sorted_by_exclusive_start_time_) { + slice_details_sorted_by_offset.push_back(&slice_detail); + } + absl::c_sort(slice_details_sorted_by_offset, [](const SliceDetail* lhs, + const SliceDetail* rhs) { + return lhs->slice_decision.chunk.offset < rhs->slice_decision.chunk.offset; + }); + + for (int i = 0; i < num_slices; ++i) { + SliceDetail* slice_detail = slice_details_sorted_by_offset[i]; + HeapSimulator::Chunk& chunk = slice_detail->slice_decision.chunk; + const AllocatedSlice& repacked_slice_data = data.slices_sorted_by_offset[i]; + chunk = HeapSimulator::Chunk::FromOffsetSize(repacked_slice_data.offset, + chunk.size); + slice_detail->copy_start_after_time = + repacked_slice_data.inclusive_start_time - 1; + slice_detail->slice_decision.exclusive_start_time = + InclusiveToExclusiveStartTime(repacked_slice_data.inclusive_start_time); + } + + absl::c_sort(slice_details_sorted_by_exclusive_start_time_, + [](const SliceDetail& lhs, const SliceDetail& rhs) { + return std::make_tuple(lhs.copy_start_after_time, + lhs.slice_decision.chunk.offset) < + std::make_tuple(rhs.copy_start_after_time, + rhs.slice_decision.chunk.offset); + }); +} + +const std::vector& +SlicedCopyAllocation::slice_details_sorted_by_start_time() const { + return slice_details_sorted_by_exclusive_start_time_; +} + +std::vector& +SlicedCopyAllocation::mutable_slice_details_sorted_by_start_time() { + return slice_details_sorted_by_exclusive_start_time_; +} + +bool SlicedCopyAllocation::operator==(const SlicedCopyAllocation& other) const { + return this->base_is_equal(static_cast(other)) && + slice_details_sorted_by_exclusive_start_time_ == + other.slice_details_sorted_by_exclusive_start_time_ && + concat_ == other.concat_; +} + +std::string SlicedCopyAllocation::ToString() const { + std::string memory_space_str = MemorySpaceToString(memory_space()); + if (memory_space() == MemorySpace::kAlternate) { + absl::StrAppend(&memory_space_str, " (off: ", maybe_chunk()->offset, ")"); + } + return absl::StrCat( + "Sliced Copy Allocation in ", memory_space_str, + ", start_time:", start_time(), ", end_time:", end_time(), + ", first_slice_copy_start_after_time: ", + slice_details_sorted_by_start_time().front().copy_start_after_time, + ", last_slice_copy_done_before_time: ", + slice_details_sorted_by_start_time().back().copy_done_before_time, + ", uses: ", UsesToString(uses()), + ", sync_mem_op: ", sync_mem_op_ ? sync_mem_op_->name() : "none", + ", from ", prev_allocation_.ToString()); +} + +absl::Status SlicedCopyAllocation::CreateBitcastConcat( + const Shape& shape, absl::Span slices) { + CHECK(!slices.empty()); + concat_ = + slices.front()->parent()->AddInstruction(HloInstruction::CreateCustomCall( + shape, slices, + xla::memory_space_assignment::kConcatBitcastCustomCall)); + return absl::OkStatus(); +} + +std::string SlicedCopyAllocation::SliceDetail::ToString() const { + return absl::StrCat("{ slice_decision: ", slice_decision.ToString(), + ", copy_start_after_time: ", copy_start_after_time, + ", copy_done_before_time: ", copy_done_before_time, " }"); +} + +std::tuple +SliceDetailToTuple(const SlicedCopyAllocation::SliceDetail& slice_detail) { + return std::make_tuple(std::ref(slice_detail.slice_decision), + slice_detail.copy_start_after_time, + slice_detail.copy_done_before_time, + slice_detail.copy_start, slice_detail.copy_done); +} + +bool SlicedCopyAllocation::SliceDetail::operator==( + const SliceDetail& other) const { + return SliceDetailToTuple(*this) == SliceDetailToTuple(other); +} + +absl::Status SlicedCopyAllocation::SliceDetail::CreateAsyncSlice( + const Shape& original_shape, HloInstruction& producer, + HloComputation& parent) { + if (original_shape.dimensions().size() != + slice_decision.sizing.slice_params.size()) { + return FailedPrecondition( + "%s", absl::StrCat("The number of SlicedCopyAllocation parameters ", + slice_decision.sizing.slice_params.size(), + " does not match the rank ", + original_shape.dimensions().size(), + " of the tensor we are slicing.")); + } + + std::vector start_indices; + start_indices.reserve(slice_decision.sizing.slice_params.size()); + std::vector limit_indices; + limit_indices.reserve(slice_decision.sizing.slice_params.size()); + std::vector strides; + strides.reserve(slice_decision.sizing.slice_params.size()); + + for (int i = 0; i < slice_decision.sizing.slice_params.size(); ++i) { + const SliceParam& slice_param = slice_decision.sizing.slice_params[i]; + start_indices.push_back(slice_param.start_inclusive); + limit_indices.push_back(slice_param.end_exclusive); + strides.push_back(1); + const int64_t new_dim = + slice_param.end_exclusive - slice_param.start_inclusive; + if (new_dim <= 0) { + return FailedPrecondition( + "%s", absl::StrCat("SlicedCopyAllocation new dimension size is ", + new_dim, ", expected something > 0.")); + } + if (original_shape.dimensions(i) < new_dim) { + return FailedPrecondition( + "%s", + absl::StrCat("SlicedCopyAllocation sliced dimension size ", new_dim, + " is bigger than its original dimension size of ", + original_shape.dimensions(i), ".")); + } + } + + HloInstruction* slice = parent.AddInstruction( + HloInstruction::CreateSlice(slice_decision.sizing.slice_shape, &producer, + start_indices, limit_indices, strides)); + ABSL_ASSIGN_OR_RETURN(copy_done, parent.CreateAsyncInstructions( + slice, {ShapeUtil::MakeShape(S32, {})})); + copy_start = copy_done->mutable_operand(0); + + return absl::OkStatus(); +} + +bool SlicedCopyAllocation::operator==(const Allocation& other) const { + const SlicedCopyAllocation* casted_other = + dynamic_cast(&other); + return casted_other != nullptr && (*this) == (*casted_other); +} + +MirroredAllocation::MirroredAllocation(const Allocation& original_allocation, + int64_t time) + : Allocation(original_allocation.defining_position(), + original_allocation.memory_space(), + original_allocation.maybe_chunk(), + /*start_time=*/time, + /*end_time=*/time, + /*cross_program_prefetch_index=*/std::nullopt), + defining_position_(std::nullopt), + original_allocation_(original_allocation) {} + +MirroredAllocation::MirroredAllocation(HloPosition defining_position, + const Allocation& original_allocation, + int64_t start_time, int64_t end_time) + : Allocation(original_allocation.defining_position(), + original_allocation.memory_space(), + original_allocation.maybe_chunk(), + /*start_time=*/start_time, + /*end_time=*/end_time, + /*cross_program_prefetch_index=*/std::nullopt), + defining_position_(defining_position), + original_allocation_(original_allocation) {} + +HloPosition MirroredAllocation::defining_position() const { + if (defining_position_.has_value()) { + return defining_position_.value(); + } + return original_allocation_.defining_position(); +} + +std::string MirroredAllocation::ToString() const { + std::string memory_space_str = MemorySpaceToString(memory_space()); + std::optional chunk = maybe_chunk(); + if (chunk) { + absl::StrAppend(&memory_space_str, " (off: ", chunk->offset, + ", size: ", chunk->size, ")"); + } + return absl::StrCat( + "MirroredAllocation in ", memory_space_str, " defined at ", + defining_position().ToString(), ", start_time:", start_time(), + ", end_time:", end_time(), ", uses: ", UsesToString(uses()), + ", for original allocation: ", original_allocation_.ToString()); +} + +absl::Status MirroredAllocation::Process( + const BitcastSplitFn& bitcast_split_fn, const HloLiveRange& hlo_live_range, + const HloAliasAnalysis& alias_analysis) { + set_original_defining_position(defining_position()); + HloInstruction* producing_instruction = AddGetTupleElements(); + HloComputation* computation = producing_instruction->parent(); + return UpdateUses(computation, producing_instruction, bitcast_split_fn, + hlo_live_range, alias_analysis); +} + +void MirroredAllocation::MarkIfNeeded( + absl::flat_hash_set& needed_allocations) const { + MarkNeeded(needed_allocations); +} + +void MirroredAllocation::MarkNeeded( + absl::flat_hash_set& needed_allocations) const { + if (needed_allocations.insert(this).second) { + original_allocation_.MarkNeeded(needed_allocations); + } +} + +bool MirroredAllocation::operator==(const Allocation& other) const { + const MirroredAllocation* casted_other = + dynamic_cast(&other); + return casted_other != nullptr && (*this) == (*casted_other) && + defining_position() == other.defining_position(); +} + +ParentAllocation::ParentAllocation(const Allocation& original_allocation, + HloInstruction* calling_instruction, + HloPosition position, int64_t time) + : Allocation(std::move(position), MemorySpace::kDefault, + original_allocation.maybe_chunk(), + /*start_time=*/time, + /*end_time=*/time, + /*cross_program_prefetch_index=*/std::nullopt), + original_allocation_(original_allocation), + calling_instruction_(calling_instruction) {} + +HloPosition ParentAllocation::defining_position() const { + return original_defining_position(); +} + +std::string ParentAllocation::ToString() const { + return absl::StrCat("Parent Allocation mirrored at ", + original_defining_position().ToString(), ", originally ", + original_allocation_.ToString()); +} + +absl::Status ParentAllocation::Process(const BitcastSplitFn& bitcast_split_fn, + const HloLiveRange& hlo_live_range, + const HloAliasAnalysis& alias_analysis) { + // Add an additional parameter to the while HLO with a reference to the buffer + // in the default memory space. + HloInstruction* producing_instruction = + original_allocation_.AddGetTupleElements(); + int new_tuple_index = calling_instruction_->shape().tuple_shapes().size(); + + HloInstruction* old_while_operand = calling_instruction_->mutable_operand(0); + ABSL_ASSIGN_OR_RETURN( + HloInstruction * new_while_operand, + TupleUtil::ReplaceTupleWith(producing_instruction, old_while_operand, + {new_tuple_index})); + AppendToTupleOriginalValue(old_while_operand, new_while_operand, + producing_instruction); + // Also replace the while op with a tuple that has the old shape. Note that we + // need to first take a snapshot of the users before calling ExtractPrefix + // since ExtractPrefix introduces additional gte users. + std::vector while_users = calling_instruction_->users(); + HloInstruction* tuple_with_old_shape = + TupleUtil::ExtractPrefix(calling_instruction_, new_tuple_index); + if (std::shared_ptr src_ov = + calling_instruction_->original_value(); + src_ov != nullptr) { + tuple_with_old_shape->CopyOriginalValue(calling_instruction_, + /*clone=*/true); + const auto& src_tree = src_ov->tree(); + for (int64_t i = 0; i < new_tuple_index; ++i) { + HloInstruction* gte = tuple_with_old_shape->mutable_operand(i); + if (src_tree.find({i}) != src_tree.end()) { + auto gte_ov = std::make_shared(gte->shape()); + CHECK_OK(gte_ov->mutable_tree()->CopyCompatibleSubtreeFrom(src_tree, + {i}, {})); + if (!gte_ov->IsEmpty()) { + gte->set_original_value(std::move(gte_ov)); + } + } + } + } + ABSL_RETURN_IF_ERROR(calling_instruction_->ReplaceOperandWithDifferentShape( + 0, new_while_operand)); + *calling_instruction_->mutable_shape() = new_while_operand->shape(); + *calling_instruction_->while_condition() + ->parameter_instruction(0) + ->mutable_shape() = new_while_operand->shape(); + *calling_instruction_->while_body() + ->parameter_instruction(0) + ->mutable_shape() = new_while_operand->shape(); + ABSL_RETURN_IF_ERROR(calling_instruction_->ReplaceAllUsesWithDifferentShape( + while_users, tuple_with_old_shape)); + + AppendToWhileLoopOriginalValue(calling_instruction_, {producing_instruction}); + HloPosition defining_position = original_defining_position(); + defining_position.index = {new_tuple_index}; + set_original_defining_position(defining_position); + + HloInstruction* final_instruction = AddGetTupleElements(); + HloComputation* computation = final_instruction->parent(); + return UpdateUses(computation, final_instruction, bitcast_split_fn, + hlo_live_range, alias_analysis); +} + +absl::Status ParentAllocation::PostProcess() { + // Update the root of the while body with the new parameter. The reason why we + // need a separate post-process for this is because other allocations may have + // while body root as a use, so they would update the old root instead of the + // new root. Doing the post-process step later ensures the root has been + // updated with other changes, and we can safely add the additional parameter. + HloComputation* while_body = calling_instruction_->while_body(); + HloInstruction* old_body_root = while_body->root_instruction(); + HloInstruction* added_element = AddGetTupleElements(); + ABSL_ASSIGN_OR_RETURN( + HloInstruction * new_while_body_root, + TupleUtil::ReplaceTupleWith(added_element, old_body_root, + original_defining_position().index)); + while_body->set_root_instruction(new_while_body_root, + /*accept_different_shape=*/true); + AppendToTupleOriginalValue(old_body_root, new_while_body_root, added_element); + return absl::OkStatus(); +} + +void ParentAllocation::MarkIfNeeded( + absl::flat_hash_set& needed_allocations) const { + // Parent allocations are only needed if they have any uses or if there is a + // copy allocation that copies this value (in that case, the copy allocation + // will call this allocation's MarkNeeded function). + if (!has_no_uses()) { + MarkNeeded(needed_allocations); + } +} + +void ParentAllocation::MarkNeeded( + absl::flat_hash_set& needed_allocations) const { + if (needed_allocations.insert(this).second) { + original_allocation_.MarkNeeded(needed_allocations); + } +} + +bool ParentAllocation::operator==(const Allocation& other) const { + const ParentAllocation* casted_other = + dynamic_cast(&other); + return casted_other != nullptr && (*this) == (*casted_other); +} + +WindowPrefetchedAllocation::WindowPrefetchedAllocation( + Allocation& prev_allocation, HloUse use, const HeapSimulator::Chunk& chunk, + int64_t prefetch_start_schedule_after_time, + int64_t prefetch_done_schedule_before_time, const Options& options) + : Allocation( + {nullptr, {}}, MemorySpace::kAlternate, chunk, + ExclusiveToInclusiveStartTime(prefetch_start_schedule_after_time), + InclusiveToExclusiveEndTime(prefetch_done_schedule_before_time), + /*cross_program_prefetch_index=*/std::nullopt), + options_(options), + defining_position_(prev_allocation.defining_position()), + use_(use), + prefetch_start_schedule_after_(prefetch_start_schedule_after_time), + prefetch_done_schedule_before_(prefetch_done_schedule_before_time), + bytes_(chunk.size) {} + +HloPosition WindowPrefetchedAllocation::defining_position() const { + if (defining_position_.instruction != nullptr) { + return defining_position_; + } + return original_defining_position(); +} + +int64_t WindowPrefetchedAllocation::earliest_available_time() const { + return prefetch_done_schedule_before_; +} + +absl::Status WindowPrefetchedAllocation::InsertWindowPrefetchInstruction( + HloInstruction* producing_instruction, HloInstruction* use_instruction, + HloComputation* computation) { + // Derive the shape for window buffer. + Shape buffer_shape = ShapeUtil::MakeShape(U8, {options_.bytes}); + Layout layout = LayoutUtil::MakeLayout({0}); + layout.set_memory_space(options_.alternate_memory_space); + *buffer_shape.mutable_layout() = layout; + // Sync flag shape + Shape sflag_shape = ShapeUtil::MakeShape(S32, {2}); + // Output shape of the WindowPrefetch op. + Shape output_shape = ShapeUtil::MakeTupleShape({buffer_shape, sflag_shape}); + + // Insert WindowPrefetch op. + HloInstruction* custom_call = + computation->AddInstruction(HloInstruction::CreateCustomCall( + output_shape, {producing_instruction}, "WindowPrefetch")); + HloInstruction* get_buffer = computation->AddInstruction( + HloInstruction::CreateGetTupleElement(buffer_shape, custom_call, 0)); + HloInstruction* get_sflag = computation->AddInstruction( + HloInstruction::CreateGetTupleElement(sflag_shape, custom_call, 1)); + use_instruction->AppendOperand(get_buffer); + use_instruction->AppendOperand(get_sflag); + + // Add backend config for the WindowPrefetch op. + int64_t parameter_index = + use_instruction->operand_index(producing_instruction); + int64_t buffer_index = use_instruction->operand_index(get_buffer); + custom_call->set_raw_backend_config_string(absl::StrCat( + "parameter_index=", parameter_index, " buffer_index=", buffer_index)); + + // The buffer's defining position is the get_tuple_element instruction. + prefetch_instruction_ = get_buffer; + return absl::OkStatus(); +} + +absl::Status WindowPrefetchedAllocation::Process( + const BitcastSplitFn& bitcast_split_fn, const HloLiveRange& hlo_live_range, + const HloAliasAnalysis& alias_analysis) { + HloInstruction* use_instruction = use_.instruction; + int64_t use_operand = use_instruction->operand_count(); + HloInstruction* producing_instruction = + use_instruction->mutable_operand(use_.operand_number); + HloComputation* computation = producing_instruction->parent(); + CHECK_EQ(use_instruction->opcode(), HloOpcode::kFusion); + + ABSL_RETURN_IF_ERROR(InsertWindowPrefetchInstruction( + producing_instruction, use_instruction, computation)); + + // Notify the backend that an operand has been appended as a window prefetch + // buffer. + options_.notify_operand_appended_fn(use_instruction, use_.operand_number, + use_operand); + + // Set the original defining position to the window prefetch instruction. + set_original_defining_position(defining_position_); + defining_position_ = {prefetch_instruction_, {}}; + AddUse(HloUse{use_instruction, use_operand}); + return absl::OkStatus(); +} + +void WindowPrefetchedAllocation::MarkIfNeeded( + absl::flat_hash_set& needed_allocations) const { + MarkNeeded(needed_allocations); +} + +void WindowPrefetchedAllocation::MarkNeeded( + absl::flat_hash_set& needed_allocations) const { + needed_allocations.insert(this); +} + +std::string WindowPrefetchedAllocation::ToString() const { + return absl::StrCat("WindowPrefetched Allocation"); +} + +bool WindowPrefetchedAllocation::operator==( + const WindowPrefetchedAllocation& other) const { + return this->base_is_equal(static_cast(other)) && + prefetch_done_schedule_before() == + other.prefetch_done_schedule_before() && + prefetch_start_schedule_after() == + other.prefetch_start_schedule_after() && + prefetch() == other.prefetch() && bytes_ == other.bytes_; +} + +bool WindowPrefetchedAllocation::operator==(const Allocation& other) const { + const WindowPrefetchedAllocation* casted_other = + dynamic_cast(&other); + return casted_other != nullptr && (*this) == (*casted_other); +} + +ScopedAllocation::ScopedAllocation(HeapSimulator::Chunk chunk, + int64_t allocation_time, + HloInstruction* defining_instruction, + bool is_post_module) + : Allocation(/*defining_position=*/{defining_instruction, {}}, + MemorySpace::kAlternate, chunk, + /*start_time=*/allocation_time, + /*end_time=*/allocation_time, + /*cross_program_prefetch_index=*/std::nullopt), + is_post_module_(is_post_module) {} + +HloPosition ScopedAllocation::defining_position() const { + return original_defining_position(); +} + +absl::Status ScopedAllocation::Process(const BitcastSplitFn& bitcast_split_fn, + const HloLiveRange& hlo_live_range, + const HloAliasAnalysis& alias_analysis) { + return absl::OkStatus(); +} + +void ScopedAllocation::MarkIfNeeded( + absl::flat_hash_set& needed_allocations) const { + MarkNeeded(needed_allocations); +} + +void ScopedAllocation::MarkNeeded( + absl::flat_hash_set& needed_allocations) const { + needed_allocations.insert(this); +} + +std::string ScopedAllocation::ToString() const { + std::string name = ""; + if (!is_post_module_) { + name = defining_position().instruction->name(); + } + return absl::StrCat("Scoped Allocation defined for ", name, + ", allocation_time: ", start_time(), + ", chunk: ", chunk().ToString()); +} + +bool ScopedAllocation::operator==(const Allocation& other) const { + const ScopedAllocation* casted_other = + dynamic_cast(&other); + return casted_other != nullptr && (*this) == (*casted_other); +} + +bool ScopedAllocation::operator==(const ScopedAllocation& other) const { + return this->base_is_equal(static_cast(other)) && + is_post_module_ == other.is_post_module_; +} + +std::tuple GetAllocationSortTuple( + const std::unique_ptr& allocation) { + int64_t scheduled_on_or_before = allocation->start_time(); + int64_t scheduled_on_or_after = allocation->start_time(); + if (allocation->is_copy_allocation()) { + auto copy_allocation = absl::down_cast(allocation.get()); + scheduled_on_or_before = copy_allocation->copy_done_schedule_before(); + scheduled_on_or_after = copy_allocation->copy_start_schedule_after(); + } + return std::forward_as_tuple(scheduled_on_or_before, + !allocation->is_copy_allocation(), + scheduled_on_or_after); +} + +void SortAllocationSequence(AllocationSequence& allocations) { + absl::c_sort(allocations, [](const std::unique_ptr& lhs, + const std::unique_ptr& rhs) { + return GetAllocationSortTuple(lhs) < GetAllocationSortTuple(rhs); + }); +} + +std::string AllocationSequenceToString(AllocationSequence& allocations, + bool sort_allocations) { + if (sort_allocations) { + SortAllocationSequence(allocations); + } + std::string allocations_str = "\n"; + for (const std::unique_ptr& allocation : allocations) { + absl::StrAppend(&allocations_str, allocation->ToString(), "\n"); + } + return allocations_str; +} + +std::vector GetAllocationSequenceInRawPointers( + AllocationSequence& allocations) { + std::vector allocations_in_raw_pointers; + for (const std::unique_ptr& allocation : allocations) { + allocations_in_raw_pointers.push_back(allocation.get()); + } + return allocations_in_raw_pointers; +} + +namespace { + +struct AllocationSummary { + static void Add(const Allocation& allocation, + std::vector& data) { + if (!allocation.is_sliced_copy_allocation()) { + data.push_back( + AllocationSummary{allocation.chunk(), allocation.start_time(), + allocation.end_time(), GetName(allocation)}); + return; + } + const SlicedCopyAllocation& sliced_copy_allocation = + dynamic_cast(allocation); + for (size_t i = 0; + i < sliced_copy_allocation.slice_details_sorted_by_start_time().size(); + ++i) { + const SlicedCopyAllocation::SliceDetail& slice_detail = + sliced_copy_allocation.slice_details_sorted_by_start_time()[i]; + data.push_back(AllocationSummary{ + slice_detail.slice_decision.chunk, + ExclusiveToInclusiveStartTime( + slice_detail.slice_decision.exclusive_start_time), + sliced_copy_allocation.end_time(), GetName(allocation, i)}); + } + } + + static std::string GetName(const Allocation& allocation, + std::optional slice_index = std::nullopt) { + std::vector attributes; + if (slice_index.has_value()) { + attributes.push_back(absl::StrCat("slice ", *slice_index)); + } + if (allocation.is_scoped_allocation()) { + attributes.push_back("scoped"); + } + if (allocation.cross_program_prefetch_index().has_value()) { + attributes.push_back("cross-program prefetch"); + } + if (attributes.empty()) { + return allocation.defining_position().ToString(); + } + return absl::StrCat(allocation.defining_position().ToString(), " (", + absl::StrJoin(attributes, ", "), ")"); + } + + HeapSimulator::Chunk chunk; + int64_t start_time_inclusive; + int64_t end_time_inclusive; + std::string name; +}; + +} // namespace + +void AllocationSequenceDebugging::LogAltMemAllocationsAt( + const AllocationSequence& allocations, int64_t time) { + std::vector data_vector; + for (const std::unique_ptr& allocation : allocations) { + if (allocation->start_time() <= time && allocation->end_time() >= time && + allocation->is_in_alternate_mem()) { + AllocationSummary::Add(*allocation, data_vector); + } + } + absl::c_sort(data_vector, + [](const AllocationSummary& a, const AllocationSummary& b) { + if (a.chunk.offset != b.chunk.offset) { + return a.chunk.offset < b.chunk.offset; + } + return a.start_time_inclusive < b.start_time_inclusive; + }); + LOG(INFO) << "Live allocations in alternate mem at instruction time " << time + << " (before MSA alters the graph):"; + for (const AllocationSummary& data : data_vector) { + LOG(INFO) << "Alt mem allocation in chunk " << data.chunk.ToString() + << " during [" << data.start_time_inclusive << "," + << data.end_time_inclusive << "], holding " << data.name; + } +} + +} // namespace xla::memory_space_assignment diff --git a/third_party/xla/xla/service/memory_space_assignment/allocation_test.cc b/third_party/xla/xla/service/memory_space_assignment/allocation_test.cc new file mode 100644 index 00000000000000..4f8339e21c078f --- /dev/null +++ b/third_party/xla/xla/service/memory_space_assignment/allocation_test.cc @@ -0,0 +1,393 @@ +/* Copyright 2024 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xla/service/memory_space_assignment/allocation.h" + +#include +#include +#include + +#include +#include +#include "absl/strings/string_view.h" +#include "xla/hlo/analysis/alias_info.h" +#include "xla/hlo/analysis/hlo_alias_analysis.h" +#include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/ir/hlo_schedule.h" +#include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" +#include "xla/hlo/utils/hlo_live_range.h" +#include "xla/service/heap_simulator/heap_simulator.h" +#include "xla/service/hlo_value.h" +#include "xla/shape_util.h" +#include "xla/xla_data.pb.h" + +namespace xla::memory_space_assignment { +namespace { + +class AllocationTest : public HloHardwareIndependentTestBase { + protected: + void RunAnalysis(HloModule* module, + const std::vector& inst_names, + std::unique_ptr& live_range, + std::unique_ptr& alias_analysis) { + HloSchedule schedule(module); + HloInstructionSequence sequence; + for (auto name : inst_names) { + sequence.push_back(FindInstruction(module, name)); + } + schedule.set_sequence(module->entry_computation(), sequence); + + AliasInfo alias_info; + ASSERT_OK_AND_ASSIGN(alias_analysis, + HloAliasAnalysis::Run(module, &alias_info)); + ASSERT_OK_AND_ASSIGN(live_range, + HloLiveRange::Run(schedule, *alias_analysis, + module->entry_computation())); + } +}; + +TEST_F(AllocationTest, CopyAllocationProcessSimple) { + absl::string_view hlo_string = R"( +HloModule module + +ENTRY entry { + p0 = f32[2,3]{1,0} parameter(0) + p1 = f32[2,3]{1,0} parameter(1) + p1_negate = f32[2,3]{1,0} negate(p1) + add = f32[2,3]{1,0} add(p0, p1_negate) + ROOT tuple = tuple(add, p0) +} + )"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + + std::unique_ptr hlo_live_range; + std::unique_ptr alias_analysis; + RunAnalysis(module.get(), {"p0", "p1", "p1_negate", "add", "tuple"}, + hlo_live_range, alias_analysis); + // HloComputation* computation = module->entry_computation(); + HloInstruction* add = FindInstruction(module.get(), "add"); + HloInstruction* p1_negate = FindInstruction(module.get(), "p1_negate"); + + HeapSimulator::Chunk p1_negate_chunk = + HeapSimulator::Chunk::FromOffsetSize(0, 24); + + PinnedAllocation p1_negate_pinned(HloPosition{p1_negate, {}}, + MemorySpace::kDefault, p1_negate_chunk, + /*start_time=*/0, + /*end_time=*/5); + CopyAllocation copy_allocation(p1_negate_pinned, MemorySpace::kAlternate, + std::nullopt, + /*copy_start_schedule_after_time=*/2, + /*copy_done_schedule_before_time=*/3, + /*end_time=*/5, std::nullopt, + /*sync_mem_op=*/nullptr); + + // Use the correct instruction and operand numbers for the add instruction + copy_allocation.AddUse(HloUse{add, 1}); // Use of p1_negate in add + BitcastSplitFn split_fn = nullptr; + ASSERT_OK( + copy_allocation.Process(split_fn, *hlo_live_range, *alias_analysis)); + + // Check copy_start and copy_done instructions. + HloInstruction* copy_start = copy_allocation.copy_start(); + ASSERT_NE(copy_start, nullptr); + EXPECT_EQ(copy_start->opcode(), HloOpcode::kCopyStart); + EXPECT_EQ(copy_start->operand(0), p1_negate); + + HloInstruction* copy_done = copy_allocation.copy_done(); + ASSERT_NE(copy_done, nullptr); + EXPECT_EQ(copy_done->opcode(), HloOpcode::kCopyDone); + EXPECT_EQ(copy_done->operand(0), copy_start); + + // Check that uses are updated. + EXPECT_EQ(add->operand(1), copy_done); + + // Check defining position + EXPECT_EQ(copy_allocation.defining_position().instruction, copy_done); +} + +TEST_F(AllocationTest, EvictedSplitShape) { + absl::string_view hlo_string = R"( +HloModule module + +ENTRY entry { + p0 = f32[2,3]{1,0} parameter(0) + p1 = f32[2,3]{1,0} parameter(1) + p1_negate = f32[2,3]{1,0:S(1)SC(0:1)} negate(p1) + add = f32[2,3]{1,0} add(p0, p1_negate) + ROOT tuple = tuple(add, p0) +} + )"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + + std::unique_ptr hlo_live_range; + std::unique_ptr alias_analysis; + RunAnalysis(module.get(), {"p0", "p1", "p1_negate", "add", "tuple"}, + hlo_live_range, alias_analysis); + // HloComputation* computation = module->entry_computation(); + HloInstruction* add = FindInstruction(module.get(), "add"); + HloInstruction* p1_negate = FindInstruction(module.get(), "p1_negate"); + + HeapSimulator::Chunk p1_negate_chunk = + HeapSimulator::Chunk::FromOffsetSize(0, 24); + + PinnedAllocation p1_negate_pinned(HloPosition{p1_negate, {}}, + MemorySpace::kAlternate, p1_negate_chunk, + /*start_time=*/0, + /*end_time=*/5); + CopyAllocation copy_allocation(p1_negate_pinned, MemorySpace::kDefault, + std::nullopt, + /*copy_start_schedule_after_time=*/2, + /*copy_done_schedule_before_time=*/3, + /*end_time=*/5, std::nullopt, + /*sync_mem_op=*/nullptr); + + // Use the correct instruction and operand numbers for the add instruction + copy_allocation.AddUse(HloUse{add, 1}); // Use of p1_negate in add + BitcastSplitFn split_fn = nullptr; + ASSERT_OK( + copy_allocation.Process(split_fn, *hlo_live_range, *alias_analysis)); + + // Check copy_start and copy_done instructions. + HloInstruction* copy_start = copy_allocation.copy_start(); + ASSERT_NE(copy_start, nullptr); + EXPECT_EQ(copy_start->opcode(), HloOpcode::kCopyStart); + EXPECT_EQ(copy_start->operand(0), p1_negate); + + HloInstruction* copy_done = copy_allocation.copy_done(); + ASSERT_NE(copy_done, nullptr); + EXPECT_EQ(copy_done->opcode(), HloOpcode::kCopyDone); + EXPECT_EQ(copy_done->operand(0), copy_start); + EXPECT_EQ(copy_done->shape().layout().split_configs().size(), 0); + + // Check that uses are updated. + EXPECT_EQ(add->operand(1), copy_done); + + // Check defining position + EXPECT_EQ(copy_allocation.defining_position().instruction, copy_done); +} + +TEST_F(AllocationTest, CopyAllocationProcessReplaceSyncSlice) { + absl::string_view hlo_string = R"( +HloModule module + +ENTRY entry { + p0 = f32[1,3]{1,0} parameter(0) + p1 = f32[2,3]{1,0} parameter(1) + p1_negate = f32[2,3]{1,0} negate(p1) + slice = f32[1,3]{1,0} slice(p1_negate), slice={[0:1], [0:3]} + add = f32[1,3]{1,0} add(p0, slice) + ROOT tuple = tuple(add, p0) +} + )"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + + std::unique_ptr hlo_live_range; + std::unique_ptr alias_analysis; + RunAnalysis(module.get(), {"p0", "p1", "p1_negate", "slice", "add", "tuple"}, + hlo_live_range, alias_analysis); + // HloComputation* computation = module->entry_computation(); + HloInstruction* add = FindInstruction(module.get(), "add"); + HloInstruction* p1_negate = FindInstruction(module.get(), "p1_negate"); + HloInstruction* slice = FindInstruction(module.get(), "slice"); + + HeapSimulator::Chunk p1_negate_chunk = + HeapSimulator::Chunk::FromOffsetSize(0, 24); + + PinnedAllocation p1_negate_pinned(HloPosition{p1_negate, {}}, + MemorySpace::kAlternate, p1_negate_chunk, + /*start_time=*/0, + /*end_time=*/5); + CopyAllocation copy_allocation(p1_negate_pinned, MemorySpace::kAlternate, + std::nullopt, + /*copy_start_schedule_after_time=*/2, + /*copy_done_schedule_before_time=*/3, + /*end_time=*/5, std::nullopt, + /*sync_mem_op=*/slice); + + // Use the correct instruction and operand numbers for the add instruction + copy_allocation.AddUse(HloUse{add, 1}); // Use of p1_negate in add + BitcastSplitFn split_fn = nullptr; + ASSERT_OK( + copy_allocation.Process(split_fn, *hlo_live_range, *alias_analysis)); + + // Check copy_start and copy_done instructions. + HloInstruction* slice_start = copy_allocation.copy_start(); + ASSERT_NE(slice_start, nullptr); + EXPECT_EQ(slice_start->opcode(), HloOpcode::kAsyncStart); + EXPECT_EQ(slice_start->operand(0), p1_negate); + + HloInstruction* slice_done = copy_allocation.copy_done(); + ASSERT_NE(slice_done, nullptr); + EXPECT_EQ(slice_done->opcode(), HloOpcode::kAsyncDone); + EXPECT_EQ(slice_done->operand(0), slice_start); + + // Check the shapes. + EXPECT_EQ(slice_done->shape(), slice->shape()); + + // Check that uses are updated. + EXPECT_EQ(add->operand(1), slice_done); + + // Check defining position + EXPECT_EQ(copy_allocation.defining_position().instruction, slice_done); +} + +TEST_F(AllocationTest, SkipTupleReconstructionForAsyncCollective) { + absl::string_view hlo_string = R"( +HloModule module + +ENTRY entry { + new_buffer = f32[2,3]{1,0} parameter(0) + cp-start = (f32[2,3]{1,0}, f32[2,3]{1,0}, u32[], u32[]) collective-permute-start(new_buffer), channel_id=1, source_target_pairs={{0,1}} + cp-done = f32[2,3]{1,0} collective-permute-done(cp-start) + ROOT tuple = tuple(cp-done) +} + )"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + + std::unique_ptr hlo_live_range; + std::unique_ptr alias_analysis; + RunAnalysis(module.get(), {"new_buffer", "cp-start", "cp-done", "tuple"}, + hlo_live_range, alias_analysis); + + HloInstruction* cp_start = FindInstruction(module.get(), "cp-start"); + HloInstruction* cp_done = FindInstruction(module.get(), "cp-done"); + HloInstruction* new_buffer = FindInstruction(module.get(), "new_buffer"); + + HeapSimulator::Chunk chunk = HeapSimulator::Chunk::FromOffsetSize(0, 24); + + PinnedAllocation pinned(HloPosition{new_buffer, {}}, MemorySpace::kAlternate, + chunk, 0, 5); + pinned.AddUse(HloUse{cp_done, 0, {0}}); + + BitcastSplitFn split_fn = nullptr; + ASSERT_OK(pinned.Process(split_fn, *hlo_live_range, *alias_analysis)); + + EXPECT_EQ(cp_done->operand(0), cp_start); +} + +TEST_F(AllocationTest, MirroredAllocationDelegatesChunkToOriginalAllocation) { + HloComputation::Builder builder("entry"); + HloInstruction* p0 = builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {2, 3}), "p0")); + PinnedAllocation original_allocation( + HloPosition{p0, {}}, MemorySpace::kAlternate, + HeapSimulator::Chunk::FromOffsetSize(-1, 64), + /*start_time=*/0, /*end_time=*/5); + MirroredAllocation mirrored_allocation(original_allocation, /*time=*/2); + original_allocation.set_offset(128); + + const Allocation& alloc = mirrored_allocation; + ASSERT_TRUE(alloc.maybe_chunk().has_value()); + EXPECT_EQ(alloc.maybe_chunk()->offset, 128); + EXPECT_EQ(alloc.chunk().offset, 128); +} + +TEST_F(AllocationTest, ParentAllocationPropagatesOriginalValue) { + absl::string_view hlo_string = R"( +HloModule module + +while_body { + param = (f32[2,3]{1,0}, f32[2,3]{1,0}) parameter(0), origin={({"bp0"}, {"bp1"})} + gte0 = f32[2,3]{1,0} get-tuple-element(param), index=0 + gte1 = f32[2,3]{1,0} get-tuple-element(param), index=1 + add = f32[2,3]{1,0} add(gte0, gte1) + ROOT root = (f32[2,3]{1,0}, f32[2,3]{1,0}) tuple(add, gte1), origin={({"br0"}, {"br1"})} +} + +while_condition { + param = (f32[2,3]{1,0}, f32[2,3]{1,0}) parameter(0), origin={({"cp0"}, {"cp1"})} + ROOT cond = pred[] constant(true) +} + +ENTRY entry { + p0 = f32[2,3]{1,0} parameter(0), origin={{"p0"}} + p1 = f32[2,3]{1,0} parameter(1), origin={{"p1"}} + p2 = f32[2,3]{1,0} parameter(2), origin={{"p2"}} + init = (f32[2,3]{1,0}, f32[2,3]{1,0}) tuple(p0, p1), origin={({"p0"}, {"p1"})} + while = (f32[2,3]{1,0}, f32[2,3]{1,0}) while(init), condition=while_condition, body=while_body, origin={({"w0"}, {"w1"}),["while#$"]} + entry_gte0 = f32[2,3]{1,0} get-tuple-element(while), index=0 + ROOT out = (f32[2,3]{1,0}) tuple(entry_gte0) +} + )"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + + std::unique_ptr hlo_live_range; + std::unique_ptr alias_analysis; + RunAnalysis(module.get(), + {"p0", "p1", "p2", "init", "while", "entry_gte0", "out"}, + hlo_live_range, alias_analysis); + + HloInstruction* p2 = FindInstruction(module.get(), "p2"); + HloInstruction* while_instr = FindInstruction(module.get(), "while"); + HloInstruction* entry_gte0 = FindInstruction(module.get(), "entry_gte0"); + + HeapSimulator::Chunk chunk = HeapSimulator::Chunk::FromOffsetSize(0, 24); + PinnedAllocation p2_pinned(HloPosition{p2, {}}, MemorySpace::kDefault, chunk, + /*start_time=*/0, /*end_time=*/10); + + ParentAllocation parent_alloc( + p2_pinned, while_instr, + HloPosition{while_instr->while_body()->parameter_instruction(0), {}}, + /*time=*/5); + + BitcastSplitFn split_fn = nullptr; + ASSERT_OK(parent_alloc.Process(split_fn, *hlo_live_range, *alias_analysis)); + ASSERT_OK(parent_alloc.PostProcess()); + + ASSERT_NE(while_instr->while_init()->original_value(), nullptr); + EXPECT_EQ(while_instr->while_init()->original_value()->ToString(), + R"(({"p0"}, {"p1"}, {"p2"}))"); + + ASSERT_NE(while_instr->original_value(), nullptr); + EXPECT_EQ(while_instr->original_value()->ToString(), + R"(({"w0"}, {"w1"}, {"p2"}),["while#$"])"); + + ASSERT_NE( + while_instr->while_body()->parameter_instruction(0)->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_body() + ->parameter_instruction(0) + ->original_value() + ->ToString(), + R"(({"bp0"}, {"bp1"}, {"p2"}))"); + + ASSERT_NE(while_instr->while_condition() + ->parameter_instruction(0) + ->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_condition() + ->parameter_instruction(0) + ->original_value() + ->ToString(), + R"(({"cp0"}, {"cp1"}, {"p2"}))"); + + ASSERT_NE(while_instr->while_body()->root_instruction()->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_body() + ->root_instruction() + ->original_value() + ->ToString(), + R"(({"br0"}, {"br1"}, {"p2"}))"); + + const HloInstruction* tuple_with_old_shape = entry_gte0->operand(0); + ASSERT_NE(tuple_with_old_shape->original_value(), nullptr); + EXPECT_EQ(tuple_with_old_shape->original_value()->ToString(), + R"(({"w0"}, {"w1"}),["while#$"])"); +} + +} // namespace +} // namespace xla::memory_space_assignment diff --git a/third_party/xla/xla/service/platform_util.cc b/third_party/xla/xla/service/platform_util.cc index 19bed4a689a89b..a1b1866ae08d87 100644 --- a/third_party/xla/xla/service/platform_util.cc +++ b/third_party/xla/xla/service/platform_util.cc @@ -29,6 +29,7 @@ limitations under the License. #include "absl/status/statusor.h" #include "absl/strings/ascii.h" #include "absl/strings/match.h" +#include "absl/strings/str_cat.h" #include "absl/strings/str_join.h" #include "absl/strings/string_view.h" #include "xla/debug_options_flags.h" @@ -271,6 +272,7 @@ PlatformUtil::GetStreamExecutors( } std::vector out; + std::vector device_errors; for (int i = 0; i < executors.size(); ++i) { absl::StatusOr& se = executors[i]; if (se.ok()) { @@ -279,11 +281,13 @@ PlatformUtil::GetStreamExecutors( LOG(ERROR) << "Failed to create stream executor for device " << platform->Name() << ":" << device_ordinals[i] << ": " << se.status().message(); + device_errors.push_back(absl::StrCat("Device ", device_ordinals[i], ": ", + se.status().message())); } } if (out.empty()) { - return Internal("no supported devices found for platform %s", - platform->Name()); + return Internal("no supported devices found for platform %s: %s", + platform->Name(), absl::StrJoin(device_errors, ", ")); } return out; } diff --git a/third_party/xla/xla/service/platform_util_test.cc b/third_party/xla/xla/service/platform_util_test.cc index 78bb47c211afc6..b0ef13178f18d0 100644 --- a/third_party/xla/xla/service/platform_util_test.cc +++ b/third_party/xla/xla/service/platform_util_test.cc @@ -15,12 +15,15 @@ limitations under the License. #include "xla/service/platform_util.h" +#include + #include #include #include "absl/status/status.h" #include "absl/status/status_matchers.h" #include "xla/stream_executor/cuda/cuda_platform_id.h" #include "xla/stream_executor/host/host_platform_id.h" +#include "xla/stream_executor/mock_platform.h" #include "xla/stream_executor/rocm/rocm_platform_id.h" #include "xla/stream_executor/sycl/sycl_platform_id.h" @@ -29,6 +32,10 @@ namespace { using ::absl_testing::IsOkAndHolds; using ::absl_testing::StatusIs; +using ::testing::HasSubstr; +using ::testing::NiceMock; +using ::testing::Return; +using ::testing::ReturnRefOfCopy; TEST(PlatformUtilTest, GetPlatformIdFromCanonicalName) { EXPECT_THAT(PlatformUtil::GetPlatformIdFromCanonicalName("host"), @@ -55,5 +62,23 @@ TEST(PlatformUtilTest, GetPlatformIdFromCanonicalName) { StatusIs(absl::StatusCode::kInvalidArgument)); } +TEST(PlatformUtilTest, GetStreamExecutorsReportsAllDeviceErrors) { + NiceMock platform; + ON_CALL(platform, id()) + .WillByDefault(Return(stream_executor::cuda::kCudaPlatformId)); + ON_CALL(platform, Name()).WillByDefault(ReturnRefOfCopy(std::string("CUDA"))); + ON_CALL(platform, VisibleDeviceCount()).WillByDefault(Return(2)); + ON_CALL(platform, ExecutorForDevice(0)) + .WillByDefault(Return(absl::InternalError("out of memory"))); + ON_CALL(platform, ExecutorForDevice(1)) + .WillByDefault(Return(absl::UnavailableError("device busy"))); + + EXPECT_THAT( + PlatformUtil::GetStreamExecutors(&platform), + StatusIs(absl::StatusCode::kInternal, + HasSubstr("no supported devices found for platform CUDA: " + "Device 0: out of memory, Device 1: device busy"))); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/service/propagate_original_value_test.cc b/third_party/xla/xla/service/propagate_original_value_test.cc index 7dc0b8c8cdfbfd..124026e64b664d 100644 --- a/third_party/xla/xla/service/propagate_original_value_test.cc +++ b/third_party/xla/xla/service/propagate_original_value_test.cc @@ -222,7 +222,7 @@ TEST_F(OriginalValueRecoveryTableTest, constexpr absl::string_view hlo_string = R"hlo( // CHECK-NOT: origin_recovery_table // CHECK: ENTRY %[[COMPUTATION:.*]] (param: f32[4]) -> f32[4] -// CHECK: parameter(0), sharding={replicated}, origin={{[{]}}{"param_origin"} +// CHECK: parameter(0), origin={{[{]}}{"param_origin"} // CHECK: negate(%param), origin={{[{]}}{"negate_origin"} HloModule test, entry_computation_layout={(f32[4]{0})->f32[4]{0}}, num_partitions=2 diff --git a/third_party/xla/xla/service/shape_inference.cc b/third_party/xla/xla/service/shape_inference.cc index 86ed2b98e42859..ed0006136a6cd3 100644 --- a/third_party/xla/xla/service/shape_inference.cc +++ b/third_party/xla/xla/service/shape_inference.cc @@ -3706,6 +3706,39 @@ ShapeInference::InferCollectivePermuteDoneShape(const Shape& operand_shape) { return operand_shape; } +/*static */ absl::StatusOr ShapeInference::InferShuffleShape( + const Shape& operand_shape, absl::Span dimensions, + const ShuffleMode& mode) { + ABSL_RETURN_IF_ERROR(ExpectArray(operand_shape, "operand of shuffle")); + if (dimensions.empty()) { + return InvalidArgument("A shuffle must shuffle at least one dimension."); + } + if (!AllUnique(dimensions)) { + return InvalidArgument("A dimension number is duplicated in shuffle."); + } + for (int64_t dimension : dimensions) { + if (dimension < 0 || dimension >= operand_shape.dimensions().size()) { + return InvalidArgument( + "One of the shuffle dimensions (%d) is out-of-bounds in shape %s.", + dimension, ShapeUtil::HumanString(operand_shape)); + } + } + // Constraints on the attributes that are specific to the shuffle mode. + switch (mode.mode_case()) { + case ShuffleMode::kRotate: + if (dimensions.size() != mode.rotate().shifts().size()) { + return InvalidArgument( + "dimensions and shifts must have the same size, got %d and %d.", + dimensions.size(), mode.rotate().shifts().size()); + } + break; + case ShuffleMode::MODE_NOT_SET: + return InvalidArgument("A shuffle must specify a mode."); + } + // A shuffle only moves elements around, so the shape is preserved. + return operand_shape; +} + /* static */ absl::StatusOr ShapeInference::InferGetTupleElementShape( const Shape& arg, int64_t index) { if (!arg.IsTuple()) { diff --git a/third_party/xla/xla/service/shape_inference.h b/third_party/xla/xla/service/shape_inference.h index 15e71e840a09c2..c5566066296781 100644 --- a/third_party/xla/xla/service/shape_inference.h +++ b/third_party/xla/xla/service/shape_inference.h @@ -248,6 +248,12 @@ class ShapeInference { static absl::StatusOr InferReverseShape( const Shape& operand_shape, absl::Span dimensions); + // Infers the shape of a shuffle instruction, which rearranges the operand's + // elements along the given dimensions according to `mode`. + static absl::StatusOr InferShuffleShape( + const Shape& operand_shape, absl::Span dimensions, + const ShuffleMode& mode); + // Infers the shape produced by a slice operation spanning from the starts to // the limits in the original shape's dimensions. // diff --git a/third_party/xla/xla/service/shape_inference_test.cc b/third_party/xla/xla/service/shape_inference_test.cc index 0db74cc0433cf8..c8c75fa2673d26 100644 --- a/third_party/xla/xla/service/shape_inference_test.cc +++ b/third_party/xla/xla/service/shape_inference_test.cc @@ -28,6 +28,7 @@ limitations under the License. #include "absl/algorithm/container.h" #include "absl/log/check.h" #include "absl/log/log.h" +#include "absl/status/status.h" #include "absl/status/status_matchers.h" #include "absl/strings/str_join.h" #include "absl/strings/string_view.h" @@ -42,6 +43,7 @@ limitations under the License. #include "xla/service/hlo.pb.h" #include "xla/shape.h" #include "xla/shape_util.h" +#include "xla/shuffle.h" #include "xla/xla_data.pb.h" #include "tsl/platform/errors.h" @@ -3073,6 +3075,57 @@ TEST_F(ShapeInferenceTest, ReverseInvalidDimension) { HasSubstr("Expected array argument")); } +TEST_F(ShapeInferenceTest, ShuffleRotate) { + const Shape input_shape = ShapeUtil::MakeShape(F32, {10, 25}); + + ASSERT_THAT(ShapeInference::InferShuffleShape(input_shape, {0, 1}, + shuffle::Rotate({2, 5})), + absl_testing::IsOkAndHolds(input_shape)); +} + +TEST_F(ShapeInferenceTest, ShuffleInvalidDimension) { + const Shape input_shape = ShapeUtil::MakeShape(F32, {10, 25}); + + ASSERT_THAT(ShapeInference::InferShuffleShape(input_shape, {0, 2}, + shuffle::Rotate({2, 5})), + absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("out-of-bounds"))); + + ASSERT_THAT(ShapeInference::InferShuffleShape(input_shape, {0, -1}, + shuffle::Rotate({2, 5})), + absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("out-of-bounds"))); + + ASSERT_THAT(ShapeInference::InferShuffleShape(input_shape, {0, 0}, + shuffle::Rotate({2, 5})), + absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("duplicated"))); + + ASSERT_THAT(ShapeInference::InferShuffleShape(input_shape, {0, 1}, + shuffle::Rotate({2})), + absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("dimensions and shifts must " + "have the same size"))); +} + +TEST_F(ShapeInferenceTest, ShuffleWithoutMode) { + const Shape input_shape = ShapeUtil::MakeShape(F32, {10, 25}); + + ASSERT_THAT( + ShapeInference::InferShuffleShape(input_shape, {0, 1}, ShuffleMode()), + absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("must specify a mode"))); +} + +TEST_F(ShapeInferenceTest, ShuffleWithoutDimensions) { + const Shape input_shape = ShapeUtil::MakeShape(F32, {10, 25}); + + ASSERT_THAT( + ShapeInference::InferShuffleShape(input_shape, {}, shuffle::Rotate({})), + absl_testing::StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("must shuffle at least one dimension"))); +} + TEST_F(ShapeInferenceTest, Call) { const absl::StatusOr inferred_shape0 = ShapeInference::InferCallShape({}, ShapeUtil::MakeProgramShape({}, f32_)); diff --git a/third_party/xla/xla/service/sharding_propagation.cc b/third_party/xla/xla/service/sharding_propagation.cc index 26e6d65e72ae87..e4a1749f858baa 100644 --- a/third_party/xla/xla/service/sharding_propagation.cc +++ b/third_party/xla/xla/service/sharding_propagation.cc @@ -207,6 +207,7 @@ const HloInstruction* PickRepresentativeOperand( case HloOpcode::kOptimizationBarrier: case HloOpcode::kRaggedAllToAll: case HloOpcode::kReverse: + case HloOpcode::kShuffle: case HloOpcode::kSlice: case HloOpcode::kShiftLeft: case HloOpcode::kShiftRightArithmetic: diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner.cc b/third_party/xla/xla/service/spmd/spmd_partitioner.cc index a866832ef944b1..4974a9b69ae604 100644 --- a/third_party/xla/xla/service/spmd/spmd_partitioner.cc +++ b/third_party/xla/xla/service/spmd/spmd_partitioner.cc @@ -232,10 +232,6 @@ bool ShouldKeepSharding(const HloInstruction* hlo, DynCast(hlo) != nullptr) { return true; } - if (hlo->opcode() == HloOpcode::kParameter && - hlo->parent() == hlo->GetModule()->entry_computation()) { - return true; - } if (keep_valid_shardings && hlo->has_sharding()) { // SPMD partitioner can generate invalid shardings since sharding is // meaningless after this pass. We only keep valid shardings to avoid @@ -253,13 +249,6 @@ absl::Status ClearShardingAttributes( auto has_unreduced_axes = [](const HloInstruction* hlo) -> bool { return hlo->frontend_attributes().map().contains(sdy::kHasUnreducedAxes); }; - for (int64_t i = 0; i < module->entry_computation()->num_parameters(); ++i) { - // Recover the unreduced sharding for parameters. - auto param = module->entry_computation()->parameter_instruction(i); - if (has_unreduced_axes(param)) { - param->set_sharding(module->spmd_parameters_shardings()[i]); - } - } const bool keep_shardings_after_spmd = module->config().debug_options().xla_keep_shardings_after_spmd(); for (HloComputation* computation : module->computations(execution_threads)) { @@ -7408,13 +7397,6 @@ absl::StatusOr SpmdPartitioner::RunImpl( ABSL_RETURN_IF_ERROR(ClearShardingAttributes( module, num_replicas() * num_partitions(), execution_threads)); - auto entry_root = module->entry_computation()->root_instruction(); - if (entry_root->has_sharding()) { - HloSharding final_sharding = - ResolveReductionOpForSharding(entry_root, entry_root->sharding()); - entry_root->set_sharding(std::move(final_sharding)); - } - if (changed) { HloPassPipeline pass("spmd-cleanup"); pass.AddPass(/*remove_cross_partition_collective_ops=*/true); @@ -7423,7 +7405,6 @@ absl::StatusOr SpmdPartitioner::RunImpl( pass.AddPass(/*is_layout_sensitive=*/false); ABSL_RETURN_IF_ERROR(pass.Run(module, execution_threads).status()); } - return changed; } diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc index 024a14599209d8..c9deeba34249bc 100644 --- a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc +++ b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc @@ -11034,10 +11034,12 @@ ENTRY entry { EXPECT_THAT(module->entry_computation()->root_instruction(), op::Tuple(sum_full, sum_manual)); - EXPECT_THAT(module->entry_computation()->parameter_instruction(0), - op::Sharding("{{devices=[2,1]<=[2]},{manual},{manual}}")); - EXPECT_THAT(module->entry_computation()->parameter_instruction(1), - op::Sharding("{{manual},{devices=[2,1]<=[2]},{manual}}")); + EXPECT_THAT(module->entry_computation()->parameter_instructions(), + Each(op::NoSharding())); + EXPECT_THAT( + module->spmd_parameters_shardings(), + ElementsAre(*ParseSharding("{{devices=[2,1]<=[2]},{manual},{manual}}"), + *ParseSharding("{{manual},{devices=[2,1]<=[2]},{manual}}"))); } TEST_P(SpmdPartitioningTest, NestedManual) { @@ -18471,9 +18473,12 @@ ENTRY entry { ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string, /*num_devices=*/2)); VLOG(1) << module->ToString(); - // Check that unreduced HloSharding is preserved after the pass. + // Check that unreduced HloSharding is preserved on the module and cleared on + // the parameter instructions after the pass. EXPECT_THAT(module->entry_computation()->parameter_instructions(), - Each(op::Sharding("{unreduced}"))); + Each(op::NoSharding())); + EXPECT_THAT(module->spmd_parameters_shardings(), + Each(*ParseSharding("{unreduced}"))); } TEST_P(SpmdPartitioningTest, SubgroupUnreducedParam) { @@ -18488,10 +18493,13 @@ ENTRY entry { ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string, /*num_devices=*/4)); VLOG(1) << module->ToString(); - // Check that unreduced HloSharding is preserved after the pass. - EXPECT_THAT( - module->entry_computation()->parameter_instructions(), - Each(op::Sharding("{devices=[1,2,2]<=[4] last_tile_dims={unreduced}}"))); + // Check that unreduced HloSharding is preserved on the module and cleared on + // the parameter instructions after the pass. + EXPECT_THAT(module->entry_computation()->parameter_instructions(), + Each(op::NoSharding())); + EXPECT_THAT(module->spmd_parameters_shardings(), + Each(*ParseSharding( + "{devices=[1,2,2]<=[4] last_tile_dims={unreduced}}"))); } TEST_P(SpmdPartitioningTest, SubgroupUnreducedDot) { @@ -19188,9 +19196,12 @@ ENTRY entry { })"; ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string, /*num_devices=*/2)); - // Check that unreduced HloSharding is preserved after the pass. - for (auto* param : module->entry_computation()->parameter_instructions()) { - EXPECT_THAT(param->sharding().ToString(), HasSubstr("unreduced")); + // Check that unreduced HloSharding is preserved on the module and cleared on + // the parameter instructions after the pass. + EXPECT_THAT(module->entry_computation()->parameter_instructions(), + Each(op::NoSharding())); + for (const auto& sharding : module->spmd_parameters_shardings()) { + EXPECT_THAT(sharding.ToString(), HasSubstr("unreduced")); } } @@ -19204,10 +19215,13 @@ ENTRY entry { })"; ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string, /*num_devices=*/4)); - // Check that unreduced HloSharding is preserved after the pass. + // Check that unreduced HloSharding is preserved on the module and cleared on + // the parameter instructions after the pass. + EXPECT_THAT(module->entry_computation()->parameter_instructions(), + Each(op::NoSharding())); EXPECT_THAT( - module->entry_computation()->parameter_instructions(), - Each(op::Sharding("{mesh['x'=2,'y'=2] [{?},{'x'}], unreduced={'y'}}"))); + module->spmd_parameters_shardings(), + Each(*ParseSharding("{mesh['x'=2,'y'=2] [{?},{'x'}], unreduced={'y'}}"))); } TEST_F(SpmdPartitioningV3Test, TupleSubgroupUnreducedParamV3) { @@ -19225,8 +19239,10 @@ ENTRY entry { ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string, /*num_devices=*/2)); EXPECT_THAT(module->entry_computation()->parameter_instructions(), - Each(op::Sharding("{{mesh['x'=2] [], unreduced={'x'}}, " - "{mesh['x'=2] [], unreduced={'x'}}}"))); + Each(op::NoSharding())); + EXPECT_THAT(module->spmd_parameters_shardings(), + Each(*ParseSharding("{{mesh['x'=2] [], unreduced={'x'}}, " + "{mesh['x'=2] [], unreduced={'x'}}}"))); } TEST_F(SpmdPartitioningV3Test, PatternMatchMergeNamedSharding) { @@ -19490,10 +19506,13 @@ ENTRY entry { ASSERT_OK_AND_ASSIGN(auto module, PartitionComputation(hlo_string, /*num_devices=*/4)); VLOG(1) << module->ToString(); - // Check that unreduced HloSharding is preserved after the pass. + // Check that unreduced HloSharding is preserved on the module and cleared on + // the parameter instructions after the pass. + EXPECT_THAT(module->entry_computation()->parameter_instructions(), + Each(op::NoSharding())); EXPECT_THAT( - module->entry_computation()->parameter_instructions(), - Each(op::Sharding("{mesh['a'=2,'b'=2], [], unreduced=max{'a','b'}}"))); + module->spmd_parameters_shardings(), + Each(*ParseSharding("{mesh['a'=2,'b'=2], [], unreduced=max{'a','b'}}"))); } // Verifies that a true 1D scatter-conflict on a reduction dimension (e.g., diff --git a/third_party/xla/xla/service/while_loop_simplifier.cc b/third_party/xla/xla/service/while_loop_simplifier.cc index 74c9615b2c254c..aa1cb7b60455cb 100644 --- a/third_party/xla/xla/service/while_loop_simplifier.cc +++ b/third_party/xla/xla/service/while_loop_simplifier.cc @@ -278,6 +278,15 @@ static absl::StatusOr RemoveDeadTupleIndices( CopyOriginalValue(while_init, new_while_init, old_to_new_tuple_idx); CopyOriginalValue(while_op, new_while_op, old_to_new_tuple_idx); + CopyOriginalValue(while_body->parameter_instruction(0), + new_while_op->while_body()->parameter_instruction(0), + old_to_new_tuple_idx); + CopyOriginalValue(while_cond->parameter_instruction(0), + new_while_op->while_condition()->parameter_instruction(0), + old_to_new_tuple_idx); + CopyOriginalValue(while_body_root, + new_while_op->while_body()->root_instruction(), + old_to_new_tuple_idx); // Create a tuple op that recreates the output of the old while op. That is, // we transform to @@ -846,7 +855,9 @@ static absl::StatusOr TryRemoveConstantParams(HloInstruction* while_op) { }; auto add_constant_elems = - [&](HloInstruction* instr) -> std::unique_ptr { + [&](HloInstruction* instr, + const HloInstruction* old_instr = + nullptr) -> std::unique_ptr { CHECK(ShapeUtil::Compatible(instr->shape(), new_while_shape)); std::vector tuple_elems; @@ -861,7 +872,11 @@ static absl::StatusOr TryRemoveConstantParams(HloInstruction* while_op) { ++j; } } - return HloInstruction::CreateTuple(tuple_elems); + auto tuple = HloInstruction::CreateTuple(tuple_elems); + if (old_instr != nullptr) { + tuple->CopyOriginalValue(old_instr); + } + return tuple; }; // Special case: constant_tuple_indices covers the whole while parameter, so @@ -883,17 +898,20 @@ static absl::StatusOr TryRemoveConstantParams(HloInstruction* while_op) { while_cond->CloneWithReplacementPairs({ while_cond->parameter_instruction(0), add_constant_elems(add_new_instr(HloInstruction::CreateParameter( - 0, new_while_shape, - while_cond->parameter_instruction(0)->name()))), + 0, new_while_shape, + while_cond->parameter_instruction(0)->name())), + while_cond->parameter_instruction(0)), }); std::unique_ptr new_while_body = while_body->CloneWithReplacementPairs( { while_body->parameter_instruction(0), - add_constant_elems(add_new_instr(HloInstruction::CreateParameter( - 0, new_while_shape, - while_cond->parameter_instruction(0)->name()))), + add_constant_elems( + add_new_instr(HloInstruction::CreateParameter( + 0, new_while_shape, + while_cond->parameter_instruction(0)->name())), + while_body->parameter_instruction(0)), }, { while_body->root_instruction(), @@ -909,22 +927,29 @@ static absl::StatusOr TryRemoveConstantParams(HloInstruction* while_op) { module->AddEmbeddedComputation(std::move(new_while_cond)), module->AddEmbeddedComputation(std::move(new_while_body)), add_new_instr(remove_constant_elems(while_init)))); - if (while_op->original_value()) { - auto new_ov = std::make_shared(new_while_op->shape()); - int64_t new_i = 0; - for (int i = 0; i < while_shape.tuple_shapes().size(); ++i) { - if (!constant_tuple_indices.count(i)) { - CHECK_OK(new_ov->mutable_tree()->CopyCompatibleSubtreeFrom( - while_op->original_value()->tree(), {i}, {new_i++})); - } + absl::flat_hash_map old_to_new_tuple_idx; + for (int64_t i = 0, new_i = 0; i < while_shape.tuple_shapes().size(); ++i) { + if (!constant_tuple_indices.count(i)) { + old_to_new_tuple_idx[i] = new_i++; } - new_while_op->set_original_value(new_ov); } + CopyOriginalValue(while_init, new_while_op->mutable_operand(0), + old_to_new_tuple_idx); + CopyOriginalValue(while_op, new_while_op, old_to_new_tuple_idx); + CopyOriginalValue(while_body->parameter_instruction(0), + new_while_op->while_body()->parameter_instruction(0), + old_to_new_tuple_idx); + CopyOriginalValue(while_cond->parameter_instruction(0), + new_while_op->while_condition()->parameter_instruction(0), + old_to_new_tuple_idx); + CopyOriginalValue(while_body->root_instruction(), + new_while_op->while_body()->root_instruction(), + old_to_new_tuple_idx); new_while_op->CopyBackendConfigFrom(while_op); CopyFrontendAttributes(while_op, new_while_op); CopyMetadata(while_op, new_while_op); ABSL_RETURN_IF_ERROR(computation->ReplaceWithNewInstruction( - while_op, add_constant_elems(new_while_op))); + while_op, add_constant_elems(new_while_op, while_op))); for (auto& instr : new_instrs) { computation->AddInstruction(std::move(instr)); } @@ -1170,11 +1195,11 @@ static absl::StatusOr TryFlattenNestedTuples(HloInstruction* while_op) { if (HasDisableWhileLoopDceAttr(while_op)) { return false; } - auto flatten_original_value = [&](HloInstruction* old_instr, + auto flatten_original_value = [&](const HloInstruction* old_instr, HloInstruction* new_instr) { if (old_instr->original_value()) { - auto new_original_value = - std::make_shared(new_instr->shape()); + auto new_original_value = std::make_shared( + new_instr->shape(), old_instr->original_value()->call_hierarchy()); int64_t i = 0; for (auto& [shape_index, original_array] : old_instr->original_value()->tree().leaves()) { @@ -1223,7 +1248,8 @@ static absl::StatusOr TryFlattenNestedTuples(HloInstruction* while_op) { return new_instrs.back().get(); }; - auto nested = [&](HloInstruction* instr) { + auto nested = [&](HloInstruction* instr, + const HloInstruction* old_instr = nullptr) { std::vector gtes; const Shape& flat_shape = instr->shape(); gtes.reserve(flat_shape.tuple_shapes().size()); @@ -1236,6 +1262,9 @@ static absl::StatusOr TryFlattenNestedTuples(HloInstruction* while_op) { CHECK(ShapeUtil::Compatible(nested_instr->shape(), while_shape)) << ShapeUtil::HumanString(nested_instr->shape()) << " vs " << ShapeUtil::HumanString(while_shape); + if (old_instr != nullptr) { + nested_instr->CopyOriginalValue(old_instr); + } return nested_instr; }; @@ -1249,8 +1278,9 @@ static absl::StatusOr TryFlattenNestedTuples(HloInstruction* while_op) { while_cond->CloneWithReplacementPairs({ while_cond->parameter_instruction(0), nested(add_new_instr(HloInstruction::CreateParameter( - 0, flattened_shape, - while_cond->parameter_instruction(0)->name()))), + 0, flattened_shape, + while_cond->parameter_instruction(0)->name())), + while_cond->parameter_instruction(0)), }); // Create a new while-body computation, where parameter 0 has a flat shape and @@ -1261,8 +1291,9 @@ static absl::StatusOr TryFlattenNestedTuples(HloInstruction* while_op) { { while_body->parameter_instruction(0), nested(add_new_instr(HloInstruction::CreateParameter( - 0, flattened_shape, - while_body->parameter_instruction(0)->name()))), + 0, flattened_shape, + while_body->parameter_instruction(0)->name())), + while_body->parameter_instruction(0)), }, { while_body->root_instruction(), @@ -1280,14 +1311,21 @@ static absl::StatusOr TryFlattenNestedTuples(HloInstruction* while_op) { new_while_op->CopyBackendConfigFrom(while_op); CopyFrontendAttributes(while_op, new_while_op); CopyMetadata(while_op, new_while_op); - ABSL_RETURN_IF_ERROR( - computation->ReplaceWithNewInstruction(while_op, nested(new_while_op))); + flatten_original_value(while_init, new_while_op->mutable_operand(0)); + flatten_original_value(while_op, new_while_op); + flatten_original_value(while_body->parameter_instruction(0), + new_while_op->while_body()->parameter_instruction(0)); + flatten_original_value( + while_cond->parameter_instruction(0), + new_while_op->while_condition()->parameter_instruction(0)); + flatten_original_value(while_body_root, + new_while_op->while_body()->root_instruction()); + ABSL_RETURN_IF_ERROR(computation->ReplaceWithNewInstruction( + while_op, nested(new_while_op, while_op))); for (auto& instr : new_instrs) { computation->AddInstruction(std::move(instr)); } - flatten_original_value(while_init, new_while_op->mutable_operand(0)); - flatten_original_value(while_op, new_while_op); return true; } @@ -1416,7 +1454,8 @@ static absl::StatusOr TryMergeInductionVariables( // Converts `instr` into a tuple of the "old" form -- that is, to a tuple with // shape `while_body->shape()` and where the induction variables are "reified" // (i.e. they have value + * ). - auto convert_to_old_form = [&](HloInstruction* instr) { + auto convert_to_old_form = [&](HloInstruction* instr, + const HloInstruction* old_instr = nullptr) { CHECK(ShapeUtil::Compatible(instr->shape(), new_while_shape)); std::vector tuple_elems; for (int i = 0; i < while_shape.tuple_shapes().size(); ++i) { @@ -1433,7 +1472,11 @@ static absl::StatusOr TryMergeInductionVariables( // Copy the original value of the induction variable to its replacement. tuple_elems.back()->CopyOriginalValue(while_body_root->operand(i)); } - return HloInstruction::CreateTuple(tuple_elems); + auto tuple = HloInstruction::CreateTuple(tuple_elems); + if (old_instr != nullptr) { + tuple->CopyOriginalValue(old_instr); + } + return tuple; }; // Converts `root` into a tuple of the "new" form -- that is, to a tuple with @@ -1485,9 +1528,11 @@ static absl::StatusOr TryMergeInductionVariables( std::unique_ptr new_while_cond = while_cond->CloneWithReplacementPairs({ while_cond->parameter_instruction(0), - convert_to_old_form(add_new_instr(HloInstruction::CreateParameter( - 0, new_while_shape, - while_cond->parameter_instruction(0)->name()))), + convert_to_old_form( + add_new_instr(HloInstruction::CreateParameter( + 0, new_while_shape, + while_cond->parameter_instruction(0)->name())), + while_cond->parameter_instruction(0)), }); // Creating the new while body proceeds in two steps. First we convert the @@ -1502,9 +1547,11 @@ static absl::StatusOr TryMergeInductionVariables( HloComputation* temp_new_while_body = module->AddEmbeddedComputation(while_body->CloneWithReplacementPairs({ while_body->parameter_instruction(0), - convert_to_old_form(add_new_instr(HloInstruction::CreateParameter( - 0, new_while_shape, - while_body->parameter_instruction(0)->name()))), + convert_to_old_form( + add_new_instr(HloInstruction::CreateParameter( + 0, new_while_shape, + while_body->parameter_instruction(0)->name())), + while_body->parameter_instruction(0)), })); std::unique_ptr new_while_body = temp_new_while_body->CloneWithReplacementPairs({ @@ -1531,13 +1578,29 @@ static absl::StatusOr TryMergeInductionVariables( if (auto original_value = while_op->original_value()) { new_while->set_original_value(original_value); } + if (auto original_value = + while_op->while_body()->parameter_instruction(0)->original_value()) { + new_while->while_body()->parameter_instruction(0)->set_original_value( + original_value); + } + if (auto original_value = while_op->while_condition() + ->parameter_instruction(0) + ->original_value()) { + new_while->while_condition()->parameter_instruction(0)->set_original_value( + original_value); + } + if (auto original_value = + while_op->while_body()->root_instruction()->original_value()) { + new_while->while_body()->root_instruction()->set_original_value( + original_value); + } if (added_trip_counter) { - AppendToWhileLoopOriginalValue(new_while, {}); + AppendToWhileLoopOriginalValue(new_while, {nullptr}); } CopyFrontendAttributes(while_op, new_while); CopyMetadata(while_op, new_while); ABSL_RETURN_IF_ERROR(computation->ReplaceWithNewInstruction( - while_op, convert_to_old_form(new_while))); + while_op, convert_to_old_form(new_while, while_op))); for (auto& instr : new_instrs) { computation->AddInstruction(std::move(instr)); } diff --git a/third_party/xla/xla/service/while_loop_simplifier_test.cc b/third_party/xla/xla/service/while_loop_simplifier_test.cc index 60aa04044f0146..37d7f9a1a3c853 100644 --- a/third_party/xla/xla/service/while_loop_simplifier_test.cc +++ b/third_party/xla/xla/service/while_loop_simplifier_test.cc @@ -1615,15 +1615,15 @@ TEST_F(WhileLoopSimplifierTest, RemoveConstantFromLoopCarryWithOriginalValue) { const std::string hlo_string = R"( HloModule Test Body { - param = (s32[1], s32[2], s32[3]) parameter(0) + param = (s32[1], s32[2], s32[3]) parameter(0), origin={({"bp0"},{"bp1"},{"bp2"})} a = s32[1] get-tuple-element(param), index=0 a.1 = s32[1] add(a, a) b = s32[2] constant({1,1}) c = s32[3] constant({10,10,10}) - ROOT tuple = (s32[1], s32[2], s32[3]) tuple(a.1, b, c) + ROOT tuple = (s32[1], s32[2], s32[3]) tuple(a.1, b, c), origin={({"br0"},{"br1"},{"br2"})} } Cond { - param = (s32[1], s32[2], s32[3]) parameter(0) + param = (s32[1], s32[2], s32[3]) parameter(0), origin={({"cp0"},{"cp1"},{"cp2"})} a = s32[1] get-tuple-element(param), index=0 b = s32[2] get-tuple-element(param), index=1 c = s32[3] get-tuple-element(param), index=2 @@ -1646,6 +1646,30 @@ TEST_F(WhileLoopSimplifierTest, RemoveConstantFromLoopCarryWithOriginalValue) { HloInstruction* while_init = while_instr->while_init(); ASSERT_NE(while_init->original_value(), nullptr); EXPECT_EQ(while_init->original_value()->ToString(), R"(({"a"}, {"c"}))"); + ASSERT_NE( + while_instr->while_body()->parameter_instruction(0)->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_body() + ->parameter_instruction(0) + ->original_value() + ->ToString(), + R"(({"bp0"}, {"bp2"}))"); + ASSERT_NE(while_instr->while_condition() + ->parameter_instruction(0) + ->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_condition() + ->parameter_instruction(0) + ->original_value() + ->ToString(), + R"(({"cp0"}, {"cp2"}))"); + ASSERT_NE(while_instr->while_body()->root_instruction()->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_body() + ->root_instruction() + ->original_value() + ->ToString(), + R"(({"br0"}, {"br2"}))"); HloInstruction* root_instr = m->entry_computation()->root_instruction(); ASSERT_NE(root_instr->original_value(), nullptr); EXPECT_EQ(root_instr->original_value()->ToString(), @@ -1698,7 +1722,7 @@ TEST_F(WhileLoopSimplifierTest, RemoveDeadTupleIndicesWithOriginalValue) { HloModule dus %while.body (arg_tuple: (f32[3], f32[2], f32[2], f32[3], s32[])) -> (f32[3], f32[2], f32[2], f32[3], s32[]) { - %arg_tuple = (f32[3], f32[2], f32[2], f32[3], s32[]) parameter(0) + %arg_tuple = (f32[3], f32[2], f32[2], f32[3], s32[]) parameter(0), origin={({"bp0"}, {"bp1"}, {"bp2"}, {"bp3"}, {"bp4"})} %get-tuple-element.0 = f32[3] get-tuple-element(%arg_tuple), index=0 %get-tuple-element.1 = f32[2] get-tuple-element(%arg_tuple), index=1 %get-tuple-element.2 = f32[2] get-tuple-element(%arg_tuple), index=2 @@ -1710,11 +1734,11 @@ TEST_F(WhileLoopSimplifierTest, RemoveDeadTupleIndicesWithOriginalValue) { %dynamic-update-slice.0 = f32[3] dynamic-update-slice(%get-tuple-element.0, %constant.v0, s32[] %constant.1) %dynamic-update-slice.3 = f32[3] dynamic-update-slice(%get-tuple-element.3, %constant.v0, s32[] %constant.1) %add = add(s32[] %get-tuple-element.4, s32[] %constant.1) - ROOT %tuple = tuple(%dynamic-update-slice.0, %get-tuple-element.1, %get-tuple-element.2, %dynamic-update-slice.3, %add) + ROOT %tuple = tuple(%dynamic-update-slice.0, %get-tuple-element.1, %get-tuple-element.2, %dynamic-update-slice.3, %add), origin={({"br0"}, {"br1"}, {"br2"}, {"br3"}, {"br4"})} } %while.condition (arg_tuple.cond:(f32[3], f32[2], f32[2], f32[3], s32[])) -> pred[] { - %arg_tuple.cond = (f32[3], f32[2], f32[2], f32[3], s32[]) parameter(0) + %arg_tuple.cond = (f32[3], f32[2], f32[2], f32[3], s32[]) parameter(0), origin={({"cp0"}, {"cp1"}, {"cp2"}, {"cp3"}, {"cp4"})} %get-tuple-element.cond = s32[] get-tuple-element(%arg_tuple.cond), index=4 %constant.3 = s32[] constant(3) ROOT %compare = pred[] compare(s32[] %get-tuple-element.cond, s32[] %constant.3), direction=LT @@ -1744,21 +1768,45 @@ ENTRY %main (arg.0: f32[3], arg.1: f32[2]) -> (f32[3], f32[2], f32[2], f32[3]) { ASSERT_NE(while_init->original_value(), nullptr); EXPECT_EQ(while_init->original_value()->ToString(), R"(({"arg.0"}, {"arg.1"}, {"constant.0"}))"); + ASSERT_NE( + while_instr->while_body()->parameter_instruction(0)->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_body() + ->parameter_instruction(0) + ->original_value() + ->ToString(), + R"(({"bp0"}, {"bp1"}, {"bp4"}))"); + ASSERT_NE(while_instr->while_condition() + ->parameter_instruction(0) + ->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_condition() + ->parameter_instruction(0) + ->original_value() + ->ToString(), + R"(({"cp0"}, {"cp1"}, {"cp4"}))"); + ASSERT_NE(while_instr->while_body()->root_instruction()->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_body() + ->root_instruction() + ->original_value() + ->ToString(), + R"(({"br0"}, {"br1"}, {"br4"}))"); } TEST_F(WhileLoopSimplifierTest, FlattenNestedTupleWithOriginalValue) { const std::string hlo_string = R"( HloModule Test Body { - param = ((s32[1]), (s32[2], s32[3], (s32[4]))) parameter(0) + param = ((s32[1]), (s32[2], s32[3], (s32[4]))) parameter(0), origin={(({"bp0"}), ({"bp1"}, {"bp2"}, ({"bp3"})))} ta = (s32[1]) get-tuple-element(param), index=0 a = s32[1] get-tuple-element(ta), index=0 a.1 = s32[1] add(a, a) tbcd = (s32[2], s32[3], (s32[4])) get-tuple-element(param), index=1 - ROOT tuple = ((s32[1]), (s32[2], s32[3], (s32[4]))) tuple(ta, tbcd) + ROOT tuple = ((s32[1]), (s32[2], s32[3], (s32[4]))) tuple(ta, tbcd), origin={(({"br0"}), ({"br1"}, {"br2"}, ({"br3"})))} } Cond { - param = ((s32[1]), (s32[2], s32[3], (s32[4]))) parameter(0) + param = ((s32[1]), (s32[2], s32[3], (s32[4]))) parameter(0), origin={(({"cp0"}), ({"cp1"}, {"cp2"}, ({"cp3"})))} ROOT cond = pred[] constant(true) } ENTRY Loop { @@ -1773,7 +1821,7 @@ TEST_F(WhileLoopSimplifierTest, FlattenNestedTupleWithOriginalValue) { {"b"}, {"c"}, ({"d"})))} ROOT while = ((s32[1]), (s32[2], s32[3], (s32[4]))) while(init), condition=Cond, body=Body, origin={(({"while.116" {0}}), ( - {"while.116" {1}}, {"while.116" {2}}, ({"while.116" {3}})))} + {"while.116" {1}}, {"while.116" {2}}, ({"while.116" {3}}))),["while#$"]} })"; ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); @@ -1783,17 +1831,41 @@ TEST_F(WhileLoopSimplifierTest, FlattenNestedTupleWithOriginalValue) { ASSERT_NE(while_instr->original_value(), nullptr); EXPECT_EQ( while_instr->original_value()->ToString(), - R"(({"while.116" {0}}, {"while.116" {1}}, {"while.116" {2}}, {"while.116" {3}}))"); + R"(({"while.116" {0}}, {"while.116" {1}}, {"while.116" {2}}, {"while.116" {3}}),["while#$"])"); HloInstruction* while_init = while_instr->while_init(); ASSERT_NE(while_init->original_value(), nullptr); EXPECT_EQ(while_init->original_value()->ToString(), R"(({"a"}, {"b"}, {"c"}, {"d"}))"); + ASSERT_NE( + while_instr->while_body()->parameter_instruction(0)->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_body() + ->parameter_instruction(0) + ->original_value() + ->ToString(), + R"(({"bp0"}, {"bp1"}, {"bp2"}, {"bp3"}))"); + ASSERT_NE(while_instr->while_condition() + ->parameter_instruction(0) + ->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_condition() + ->parameter_instruction(0) + ->original_value() + ->ToString(), + R"(({"cp0"}, {"cp1"}, {"cp2"}, {"cp3"}))"); + ASSERT_NE(while_instr->while_body()->root_instruction()->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_body() + ->root_instruction() + ->original_value() + ->ToString(), + R"(({"br0"}, {"br1"}, {"br2"}, {"br3"}))"); } const char* const kSimpleMergeInductionVariablesModuleWithOriginalValue = R"( HloModule Test Body { - param = (TYPE[], TYPE[], TYPE[]) parameter(0) + param = (TYPE[], TYPE[], TYPE[]) parameter(0), origin={({"bp0"}, {"bp1"}, {"bp2"})} a = TYPE[] get-tuple-element(param), index=0 one = TYPE[] constant(1) @@ -1805,10 +1877,10 @@ const char* const kSimpleMergeInductionVariablesModuleWithOriginalValue = R"( c = TYPE[] add(a, b) - ROOT tuple = (TYPE[], TYPE[], TYPE[]) tuple(a1,b1,c) + ROOT tuple = (TYPE[], TYPE[], TYPE[]) tuple(a1,b1,c), origin={({"br0"}, {"br1"}, {"br2"})} } Cond { - param = (TYPE[], TYPE[], TYPE[]) parameter(0) + param = (TYPE[], TYPE[], TYPE[]) parameter(0), origin={({"cp0"}, {"cp1"}, {"cp2"})} a = TYPE[] get-tuple-element(param), index=0 b = TYPE[] get-tuple-element(param), index=1 sum = TYPE[] power(a, b) @@ -1843,6 +1915,30 @@ TEST_F(WhileLoopSimplifierTest, MergeInductionVariablesWithOriginalValue) { ASSERT_NE(while_init->original_value(), nullptr); EXPECT_EQ(while_init->original_value()->ToString(), R"(({"a"}, {"b"}, {"c"}, {}))"); + ASSERT_NE( + while_instr->while_body()->parameter_instruction(0)->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_body() + ->parameter_instruction(0) + ->original_value() + ->ToString(), + R"(({"bp0"}, {"bp1"}, {"bp2"}, {}))"); + ASSERT_NE(while_instr->while_condition() + ->parameter_instruction(0) + ->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_condition() + ->parameter_instruction(0) + ->original_value() + ->ToString(), + R"(({"cp0"}, {"cp1"}, {"cp2"}, {}))"); + ASSERT_NE(while_instr->while_body()->root_instruction()->original_value(), + nullptr); + EXPECT_EQ(while_instr->while_body() + ->root_instruction() + ->original_value() + ->ToString(), + R"(({"br0"}, {"br1"}, {"br2"}, {}))"); const HloInstruction* add = module->entry_computation()->root_instruction()->operand(0); const HloInstruction* induction_var_0 = add->operand(0); diff --git a/third_party/xla/xla/service/while_util.cc b/third_party/xla/xla/service/while_util.cc index 8e4c69ea77e656..0dc7ebd1b15889 100644 --- a/third_party/xla/xla/service/while_util.cc +++ b/third_party/xla/xla/service/while_util.cc @@ -285,40 +285,33 @@ WhileUtil::MakeInstructionsLiveIn( HloInstruction* new_while = while_instr->AddInstruction( HloInstruction::CreateWhile(new_while_shape, new_while_condition, new_while_body, new_while_init)); - if (while_instr->original_value() != nullptr) { - OriginalValue new_original_value( - new_while_shape, while_instr->original_value()->call_hierarchy()); - for (auto& [shape_index, original_array] : - new_original_value.mutable_original_arrays()) { - // The hoisted instructions are appended to the end of the while - // instruction, so the shape indices smaller than - // `elements_in_old_while_shape` are all the old original arrays that - // need to be propagated. - if (shape_index[0] < elements_in_old_while_shape) { - original_array = - while_instr->original_value()->tree().element(shape_index); - } else { - // If the element is new (i.e., hoisted), fetch its original value from - // the instruction that was hoisted. - int instruction_idx = shape_index[0] - elements_in_old_while_shape; - HloInstruction* instruction = instructions[instruction_idx]; - if (instruction->original_value() != nullptr) { - ShapeIndex shape_index_in_instruction = shape_index; - shape_index_in_instruction.erase(shape_index_in_instruction.begin()); - original_array = instruction->original_value()->original_array( - shape_index_in_instruction); - } - } - } - new_while->set_original_value( - std::make_shared(std::move(new_original_value))); + new_while_init->set_original_value( + while_instr->while_init()->original_value()); + new_while->set_original_value(while_instr->original_value()); + if (while_instr->while_body()->num_parameters() == 1) { + CHECK_EQ(new_while_body->num_parameters(), 1); + new_while_body->parameter_instruction(0)->set_original_value( + while_instr->while_body()->parameter_instruction(0)->original_value()); } + if (while_instr->while_condition()->num_parameters() == 1) { + CHECK_EQ(new_while_condition->num_parameters(), 1); + new_while_condition->parameter_instruction(0)->set_original_value( + while_instr->while_condition() + ->parameter_instruction(0) + ->original_value()); + } + new_while_body->root_instruction()->set_original_value( + while_instr->while_body()->root_instruction()->original_value()); + AppendToWhileLoopOriginalValue( + new_while, HloInstruction::InstructionVector(instructions.begin(), + instructions.end())); // We want to get rid of the old while instruction even if it has side // effecting operations so we do a manual HloComputation::RemoveInstruction // instead of relying on HloComputation::ReplaceInstruction. HloInstruction* replacement_instr = TupleUtil::ExtractPrefix( new_while, while_instr->shape().tuple_shapes().size()); + replacement_instr->CopyOriginalValue(while_instr); ABSL_RETURN_IF_ERROR(new_while->CopyAllControlDepsFrom(while_instr)); ABSL_RETURN_IF_ERROR(while_instr->DropAllControlDeps()); ABSL_RETURN_IF_ERROR(while_instr->ReplaceAllUsesWith(replacement_instr)); @@ -328,12 +321,15 @@ WhileUtil::MakeInstructionsLiveIn( std::vector live_in_instructions; for (int64_t i = elements_in_old_while_shape; i < new_while_shape.tuple_shapes().size(); i++) { - live_in_instructions.push_back(new_while_body->AddInstruction( + HloInstruction* live_in_gte = new_while_body->AddInstruction( HloInstruction::CreateGetTupleElement( instructions[i - elements_in_old_while_shape]->shape(), while_body_param, i), absl::StrCat(new_while_body->name(), ".in.", - i - elements_in_old_while_shape))); + i - elements_in_old_while_shape)); + live_in_gte->CopyOriginalValue( + instructions[i - elements_in_old_while_shape]); + live_in_instructions.push_back(live_in_gte); } WhileUtil::MakeInstructionsLiveInResult result; @@ -650,13 +646,13 @@ absl::Status WhileUtil::IncrementWhileLoopTripCount( void AppendToWhileLoopOriginalValue( HloInstruction* while_instr, const HloInstruction::InstructionVector& new_while_input_tuple_elements) { + absl::flat_hash_set updated_instrs; auto append_to_original_value = [&](HloInstruction* instr, int64_t next_index) { - std::shared_ptr old_original_value = instr->original_value(); - if (old_original_value != nullptr && - old_original_value->IsCompatibleWith(instr->shape())) { + if (instr == nullptr || !updated_instrs.insert(instr).second) { return; } + std::shared_ptr old_original_value = instr->original_value(); // Returns if neither the instruction nor any of its new tuple elements have // an original value. @@ -666,7 +662,8 @@ void AppendToWhileLoopOriginalValue( new_while_input_tuple_elements.end(), [&has_original_value](const HloInstruction* instr) { has_original_value |= - (instr->original_value() != nullptr && + (instr != nullptr && + instr->original_value() != nullptr && !instr->original_value()->IsEmpty()); }); if (!has_original_value) { @@ -674,27 +671,59 @@ void AppendToWhileLoopOriginalValue( } } + std::optional call_hierarchy; + if (old_original_value != nullptr) { + call_hierarchy = old_original_value->call_hierarchy(); + } else if (instr->opcode() == HloOpcode::kWhile) { + call_hierarchy = absl::StrCat(instr->name(), "#$"); + } + std::shared_ptr new_original_value = - std::make_shared(instr->shape()); + std::make_shared(instr->shape(), + std::move(call_hierarchy)); if (old_original_value != nullptr) { if (!old_original_value->IsTuple()) { - new_original_value->mutable_tree()->CopySubtreeFrom( - old_original_value->tree(), {}, {0}); + CHECK_EQ(next_index, 1) + << "Expected next_index == 1 when converting non-tuple while loop " + "to tuple: " + << instr->ToString(); + CHECK_OK(new_original_value->mutable_tree()->CopyCompatibleSubtreeFrom( + old_original_value->tree(), {}, {0})) + << "Incompatible OriginalValue subtree when converting non-tuple " + "while loop to tuple: " + << instr->ToString(); } else { - for (auto& [shape_index, original_array] : old_original_value->tree()) { - *new_original_value->mutable_original_array(shape_index) = - original_array; + int64_t old_tuple_size = old_original_value->tree().num_children(); + CHECK(old_tuple_size == next_index || + old_tuple_size == instr->shape().tuple_shapes().size()) + << "Unexpected old OriginalValue tuple size " << old_tuple_size + << " vs next_index " << next_index << " and widened tuple size " + << instr->shape().tuple_shapes().size() + << " for instruction: " << instr->ToString(); + for (int64_t i = 0; i < old_tuple_size; ++i) { + CHECK_OK( + new_original_value->mutable_tree()->CopyCompatibleSubtreeFrom( + old_original_value->tree(), {i}, {i})) + << "Incompatible OriginalValue subtree at index " << i << ": " + << instr->ToString(); } } } for (int64_t i = 0; i < new_while_input_tuple_elements.size(); ++i) { - if (new_while_input_tuple_elements[i]->original_value() != nullptr) { - new_original_value->mutable_tree()->CopySubtreeFrom( + if (new_while_input_tuple_elements[i] != nullptr && + new_while_input_tuple_elements[i]->original_value() != nullptr) { + CHECK_OK(new_original_value->mutable_tree()->CopyCompatibleSubtreeFrom( new_while_input_tuple_elements[i]->original_value()->tree(), {}, - {next_index + i}); + {next_index + i})) + << "Incompatible OriginalValue subtree for appended while input " + "element " + << i << ": " << new_while_input_tuple_elements[i]->ToString(); } } + CHECK(new_original_value->IsCompatibleWith(instr->shape())) + << "OriginalValue is incompatible with instruction shape: " + << instr->ToString(); instr->set_original_value(new_original_value); }; @@ -711,8 +740,20 @@ void AppendToWhileLoopOriginalValue( while_shape.tuple_shapes().size() - new_while_input_tuple_elements.size(); append_to_original_value(while_instr->while_init(), next_index); append_to_original_value(while_instr, next_index); - append_to_original_value(while_instr->while_body()->root_instruction(), - next_index); + HloInstruction* body_root = while_instr->while_body()->root_instruction(); + if (body_root->shape().IsTuple() && + body_root->shape().tuple_shapes().size() == + while_shape.tuple_shapes().size()) { + append_to_original_value(body_root, next_index); + } + if (while_instr->while_body()->num_parameters() == 1) { + append_to_original_value( + while_instr->while_body()->parameter_instruction(0), next_index); + } + if (while_instr->while_condition()->num_parameters() == 1) { + append_to_original_value( + while_instr->while_condition()->parameter_instruction(0), next_index); + } } bool WhileUtil::IsUpdatedBufferWriteOnly(const HloInstruction* instr) { diff --git a/third_party/xla/xla/service/while_util_test.cc b/third_party/xla/xla/service/while_util_test.cc index 7eefce7a2855b6..5644f9771f4f94 100644 --- a/third_party/xla/xla/service/while_util_test.cc +++ b/third_party/xla/xla/service/while_util_test.cc @@ -26,10 +26,12 @@ limitations under the License. #include "absl/strings/string_view.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/ir/hlo_original_value.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" #include "xla/hlo/testlib/test.h" #include "xla/hlo/testlib/verified_hlo_module.h" #include "xla/hlo/utils/hlo_matchers.h" +#include "xla/shape.h" #include "xla/tsl/platform/logging.h" #include "xla/util.h" @@ -503,5 +505,204 @@ ENTRY entry { EXPECT_FALSE(WhileUtil::IsUpdatedBufferWriteOnly(dus4)); } +TEST_F(WhileUtilTest, + AppendToWhileLoopOriginalValueUpdatesAllLoopInstructions) { + const char* const hlo_string = R"( +HloModule ModuleWithWhile + +while_body { + ROOT p_body = (f32[], f32[]) parameter(0), origin={({"p0"}, {"p1"})} +} + +while_condition { + p_cond = (f32[], f32[]) parameter(0), origin={({"p0"}, {"p1"})} + ROOT result = pred[] constant(true) +} + +ENTRY entry { + p0 = f32[] parameter(0), origin={{"p0"}} + p1 = f32[] parameter(1), origin={{"p1"}} + p2 = f32[] parameter(2), origin={{"p2"}} + init = (f32[], f32[]) tuple(p0, p1), origin={({"p0"}, {"p1"})} + while = (f32[], f32[]) while(init), condition=while_condition, body=while_body, origin={({"p0"}, {"p1"}),["while#$"]} + gte0 = f32[] get-tuple-element(while), index=0 + ROOT root = (f32[]) tuple(gte0) +} +)"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloInstruction* while_instr = FindInstruction(module.get(), "while"); + HloInstruction* p2 = module->entry_computation()->parameter_instruction(2); + + HloComputation* body = while_instr->while_body(); + HloComputation* cond = while_instr->while_condition(); + + Shape widened_shape = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeScalarShape(F32), ShapeUtil::MakeScalarShape(F32), + ShapeUtil::MakeScalarShape(F32)}); + *while_instr->mutable_shape() = widened_shape; + *body->parameter_instruction(0)->mutable_shape() = widened_shape; + *cond->parameter_instruction(0)->mutable_shape() = widened_shape; + *body->root_instruction()->mutable_shape() = widened_shape; + + HloInstruction* p0 = module->entry_computation()->parameter_instruction(0); + HloInstruction* p1 = module->entry_computation()->parameter_instruction(1); + HloInstruction* new_init = module->entry_computation()->AddInstruction( + HloInstruction::CreateTuple({p0, p1, p2})); + ASSERT_OK(while_instr->ReplaceOperandWithDifferentShape(0, new_init)); + + AppendToWhileLoopOriginalValue(while_instr, {p2}); + + ASSERT_NE(while_instr->original_value(), nullptr); + EXPECT_TRUE(while_instr->original_value()->IsCompatibleWith(widened_shape)); + EXPECT_THAT(while_instr->original_value()->original_array({2}), + ::testing::Optional(::testing::Eq(OriginalArray{"p2"}))); + ASSERT_TRUE(while_instr->original_value()->call_hierarchy().has_value()); + EXPECT_EQ(*while_instr->original_value()->call_hierarchy(), "while#$"); + + ASSERT_NE(body->parameter_instruction(0)->original_value(), nullptr); + EXPECT_TRUE( + body->parameter_instruction(0)->original_value()->IsCompatibleWith( + widened_shape)); + EXPECT_THAT( + body->parameter_instruction(0)->original_value()->original_array({2}), + ::testing::Optional(::testing::Eq(OriginalArray{"p2"}))); + + ASSERT_NE(cond->parameter_instruction(0)->original_value(), nullptr); + EXPECT_TRUE( + cond->parameter_instruction(0)->original_value()->IsCompatibleWith( + widened_shape)); + EXPECT_THAT( + cond->parameter_instruction(0)->original_value()->original_array({2}), + ::testing::Optional(::testing::Eq(OriginalArray{"p2"}))); + + ASSERT_NE(body->root_instruction()->original_value(), nullptr); + EXPECT_TRUE(body->root_instruction()->original_value()->IsCompatibleWith( + widened_shape)); +} + +TEST_F(WhileUtilTest, + AppendToWhileLoopOriginalValueIncompatibleSubtreeCrashes) { + const char* const hlo_string = R"( +HloModule ModuleWithWhile + +while_body { + ROOT p_body = (f32[], f32[]) parameter(0), origin={({"p0"}, {"p1"})} +} + +while_condition { + p_cond = (f32[], f32[]) parameter(0), origin={({"p0"}, {"p1"})} + ROOT result = pred[] constant(true) +} + +ENTRY entry { + p0 = f32[] parameter(0), origin={{"p0"}} + p1 = f32[] parameter(1), origin={{"p1"}} + p2 = (f32[], f32[]) parameter(2), origin={({"p2_0"}, {"p2_1"})} + init = (f32[], f32[]) tuple(p0, p1), origin={({"p0"}, {"p1"})} + while = (f32[], f32[]) while(init), condition=while_condition, body=while_body, origin={({"p0"}, {"p1"}),["while#$"]} + gte0 = f32[] get-tuple-element(while), index=0 + ROOT root = (f32[]) tuple(gte0) +} +)"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloInstruction* while_instr = FindInstruction(module.get(), "while"); + HloInstruction* p2 = module->entry_computation()->parameter_instruction(2); + + EXPECT_DEATH( + { + Shape widened_shape = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeScalarShape(F32), ShapeUtil::MakeScalarShape(F32), + ShapeUtil::MakeScalarShape(F32)}); + *while_instr->mutable_shape() = widened_shape; + AppendToWhileLoopOriginalValue(while_instr, {p2}); + }, + "Incompatible OriginalValue subtree for appended while input element 0"); +} + +TEST_F(WhileUtilTest, MakeInstructionsLiveInPropagatesOriginalValue) { + const char* const hlo_string = R"( +HloModule ModuleWithWhile + +while_body { + ROOT p_body = (f32[], f32[]) parameter(0), origin={({"p0"}, {"p1"})} +} + +while_condition { + p_cond = (f32[], f32[]) parameter(0), origin={({"p0"}, {"p1"})} + ROOT result = pred[] constant(true) +} + +ENTRY entry { + p0 = f32[] parameter(0), origin={{"p0"}} + p1 = f32[] parameter(1), origin={{"p1"}} + p2 = f32[] parameter(2), origin={{"p2"}} + init = (f32[], f32[]) tuple(p0, p1), origin={({"p0"}, {"p1"})} + while = (f32[], f32[]) while(init), condition=while_condition, body=while_body, origin={({"p0"}, {"p1"}),["while#$"]} + gte0 = f32[] get-tuple-element(while), index=0 + ROOT root = (f32[]) tuple(gte0) +} +)"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + HloInstruction* while_instr = FindInstruction(module.get(), "while"); + HloInstruction* p2 = module->entry_computation()->parameter_instruction(2); + + ASSERT_OK_AND_ASSIGN(WhileUtil::MakeInstructionsLiveInResult result, + WhileUtil::MakeInstructionsLiveIn(while_instr, {p2})); + + HloInstruction* new_while = result.new_while_instr; + ASSERT_NE(new_while->original_value(), nullptr); + EXPECT_TRUE( + new_while->original_value()->IsCompatibleWith(new_while->shape())); + EXPECT_THAT(new_while->original_value()->original_array({2}), + ::testing::Optional(::testing::Eq(OriginalArray{"p2"}))); + + ASSERT_NE(new_while->while_init()->original_value(), nullptr); + EXPECT_TRUE(new_while->while_init()->original_value()->IsCompatibleWith( + new_while->while_init()->shape())); + EXPECT_THAT(new_while->while_init()->original_value()->original_array({2}), + ::testing::Optional(::testing::Eq(OriginalArray{"p2"}))); + + ASSERT_NE(new_while->while_body()->parameter_instruction(0)->original_value(), + nullptr); + EXPECT_TRUE( + new_while->while_body() + ->parameter_instruction(0) + ->original_value() + ->IsCompatibleWith( + new_while->while_body()->parameter_instruction(0)->shape())); + EXPECT_THAT(new_while->while_body() + ->parameter_instruction(0) + ->original_value() + ->original_array({2}), + ::testing::Optional(::testing::Eq(OriginalArray{"p2"}))); + + ASSERT_NE( + new_while->while_condition()->parameter_instruction(0)->original_value(), + nullptr); + EXPECT_TRUE( + new_while->while_condition() + ->parameter_instruction(0) + ->original_value() + ->IsCompatibleWith( + new_while->while_condition()->parameter_instruction(0)->shape())); + + ASSERT_NE(new_while->while_body()->root_instruction()->original_value(), + nullptr); + EXPECT_TRUE(new_while->while_body() + ->root_instruction() + ->original_value() + ->IsCompatibleWith( + new_while->while_body()->root_instruction()->shape())); + + ASSERT_NE(result.replacement_instr->original_value(), nullptr); + EXPECT_TRUE(result.replacement_instr->original_value()->IsCompatibleWith( + result.replacement_instr->shape())); + ASSERT_EQ(result.while_body_live_in_values.size(), 1); + ASSERT_NE(result.while_body_live_in_values[0]->original_value(), nullptr); + EXPECT_THAT( + result.while_body_live_in_values[0]->original_value()->original_array({}), + ::testing::Optional(::testing::Eq(OriginalArray{"p2"}))); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/shuffle.cc b/third_party/xla/xla/shuffle.cc new file mode 100644 index 00000000000000..49c0e573a6f3ee --- /dev/null +++ b/third_party/xla/xla/shuffle.cc @@ -0,0 +1,39 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xla/shuffle.h" + +#include + +#include "absl/log/check.h" +#include "absl/types/span.h" +#include "xla/xla_data.pb.h" + +namespace xla { +namespace shuffle { + +ShuffleMode Rotate(absl::Span shifts) { + ShuffleMode mode; + mode.mutable_rotate()->mutable_shifts()->Assign(shifts.begin(), shifts.end()); + return mode; +} + +int64_t NormalizeShift(int64_t shift, int64_t dim_size) { + CHECK_GE(dim_size, 0) << "a dimension cannot have a negative size"; + return dim_size == 0 ? 0 : ((shift % dim_size) + dim_size) % dim_size; +} + +} // namespace shuffle +} // namespace xla diff --git a/third_party/xla/xla/shuffle.h b/third_party/xla/xla/shuffle.h new file mode 100644 index 00000000000000..87bb9b48c92ef5 --- /dev/null +++ b/third_party/xla/xla/shuffle.h @@ -0,0 +1,41 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef XLA_SHUFFLE_H_ +#define XLA_SHUFFLE_H_ + +#include + +#include "absl/types/span.h" +#include "xla/xla_data.pb.h" + +namespace xla { +namespace shuffle { + +// The mode of a shuffle is the pattern in which it moves the elements of its +// shuffled dimensions, and the attributes that parameterize the pattern. + +// Returns the mode of a shuffle that rotates each shuffled dimension to the +// left by the corresponding entry of `shifts`. +ShuffleMode Rotate(absl::Span shifts); + +// Returns the normalized shift in `[0, dim_size)`. Returns 0 if `dim_size` is +// 0, because every rotation of an empty dimension is a no-op. +int64_t NormalizeShift(int64_t shift, int64_t dim_size); + +} // namespace shuffle +} // namespace xla + +#endif // XLA_SHUFFLE_H_ diff --git a/third_party/xla/xla/shuffle_test.cc b/third_party/xla/xla/shuffle_test.cc new file mode 100644 index 00000000000000..151b30d1fa336b --- /dev/null +++ b/third_party/xla/xla/shuffle_test.cc @@ -0,0 +1,74 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xla/shuffle.h" + +#include +#include +#include "xla/xla_data.pb.h" + +namespace xla { +namespace shuffle { +namespace { + +using ::testing::ElementsAre; +using ::testing::IsEmpty; + +TEST(ShuffleTest, RotateSetsTheShifts) { + const ShuffleMode mode = Rotate({2, 3}); + + EXPECT_EQ(mode.mode_case(), ShuffleMode::kRotate); + EXPECT_THAT(mode.rotate().shifts(), ElementsAre(2, 3)); +} + +TEST(ShuffleTest, RotateWithoutShiftsIsStillInRotateMode) { + const ShuffleMode mode = Rotate({}); + + EXPECT_EQ(mode.mode_case(), ShuffleMode::kRotate); + EXPECT_THAT(mode.rotate().shifts(), IsEmpty()); +} + +TEST(ShuffleTest, NormalizeShiftKeepsShiftsThatAreAlreadyInRange) { + EXPECT_EQ(NormalizeShift(0, 4), 0); + EXPECT_EQ(NormalizeShift(1, 4), 1); + EXPECT_EQ(NormalizeShift(3, 4), 3); +} + +TEST(ShuffleTest, NormalizeShiftWrapsShiftsAround) { + EXPECT_EQ(NormalizeShift(4, 4), 0); + EXPECT_EQ(NormalizeShift(5, 4), 1); + EXPECT_EQ(NormalizeShift(9, 4), 1); +} + +TEST(ShuffleTest, NormalizeShiftWrapsNegativeShiftsAround) { + EXPECT_EQ(NormalizeShift(-1, 4), 3); + EXPECT_EQ(NormalizeShift(-4, 4), 0); + EXPECT_EQ(NormalizeShift(-5, 4), 3); +} + +TEST(ShuffleTest, NormalizeShiftOfADimensionOfSizeOneIsANoOp) { + EXPECT_EQ(NormalizeShift(3, 1), 0); + EXPECT_EQ(NormalizeShift(-3, 1), 0); +} + +TEST(ShuffleTest, NormalizeShiftOfAnEmptyDimensionIsANoOp) { + EXPECT_EQ(NormalizeShift(0, 0), 0); + EXPECT_EQ(NormalizeShift(3, 0), 0); + EXPECT_EQ(NormalizeShift(-3, 0), 0); +} + +} // namespace +} // namespace shuffle +} // namespace xla diff --git a/third_party/xla/xla/stream_executor/BUILD b/third_party/xla/xla/stream_executor/BUILD index fbfaaba6f928b0..2d863d1a66fc41 100644 --- a/third_party/xla/xla/stream_executor/BUILD +++ b/third_party/xla/xla/stream_executor/BUILD @@ -485,7 +485,7 @@ cc_library( hdrs = ["launch_dim.h"], visibility = internal_visibility([ ":friends", - "//xla/backends/gpu/libraries/native_custom_call_thunks:handler_allowlist", + "//xla:native_custom_call_handler_allowlist", ]), deps = [ ":launch_dim_proto_cc", @@ -874,7 +874,7 @@ cc_library( hdrs = ["kernel_spec.h"], visibility = internal_visibility([ ":friends", - "//xla/backends/gpu/libraries/native_custom_call_thunks:handler_allowlist", + "//xla:native_custom_call_handler_allowlist", ]), deps = [ ":kernel", diff --git a/third_party/xla/xla/stream_executor/command_buffer.h b/third_party/xla/xla/stream_executor/command_buffer.h index e10feee73baadc..f33039e255bacb 100644 --- a/third_party/xla/xla/stream_executor/command_buffer.h +++ b/third_party/xla/xla/stream_executor/command_buffer.h @@ -169,12 +169,28 @@ class CommandBuffer { const KernelArgs& args, absl::Span dependencies, StreamPriority priority = StreamPriority::Default) = 0; + // Creates a kernel launch command from an externally-managed native function. + // Nolint because it matches existing CreateLaunch virtual method signature + // precedent for default arguments. + // NOLINTNEXTLINE + virtual absl::StatusOr CreateLaunch( + const ThreadDim& threads, const BlockDim& blocks, + const std::optional& cluster_dims, const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args, + absl::Span dependencies, + StreamPriority priority = StreamPriority::Default) = 0; + // Updates a kernel launch command. virtual absl::Status UpdateLaunch( const Command* command, const ThreadDim& threads, const BlockDim& blocks, const std::optional& cluster_dims, const Kernel& kernel, const KernelArgs& args) = 0; + virtual absl::Status UpdateLaunch( + const Command* command, const ThreadDim& threads, const BlockDim& blocks, + const std::optional& cluster_dims, const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args) = 0; + // Type-safe wrapper for launching typed kernels. Notice that the order of // arguments is different do disambiguate from the regular launch API. template diff --git a/third_party/xla/xla/stream_executor/cuda/BUILD b/third_party/xla/xla/stream_executor/cuda/BUILD index 409f089c94351a..363a2b440b0e2b 100644 --- a/third_party/xla/xla/stream_executor/cuda/BUILD +++ b/third_party/xla/xla/stream_executor/cuda/BUILD @@ -1279,8 +1279,8 @@ xla_test( "gpu": [ "multi_gpu", ], - "gb200": ["broken"], # TODO(b/505963304): re-enable when fixed. - "gb300": ["broken"], # TODO(b/505963304): re-enable when fixed. + "gb200": ["broken"], # TODO(b/505963049): re-enable when fixed. + "gb300": ["broken"], # TODO(b/505963049): re-enable when fixed. }, backends = ["gpu"], tags = ["cuda-only"], @@ -1721,8 +1721,8 @@ xla_test( "gpu": [ "multi_gpu", ], - "gb200": ["broken"], # TODO(b/505963304): re-enable when fixed. - "gb300": ["broken"], # TODO(b/505963304): re-enable when fixed. + "gb200": ["broken"], # TODO(b/505963049): re-enable when fixed. + "gb300": ["broken"], # TODO(b/505963049): re-enable when fixed. }, backends = ["gpu"], tags = ["cuda-only"], diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.cc b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.cc index a9f5c52482a0b1..1feae5877b0801 100644 --- a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.cc +++ b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.cc @@ -722,6 +722,29 @@ absl::StatusOr CudaCommandBuffer::CreateKernelNode( const ThreadDim& threads, const BlockDim& blocks, const std::optional& cluster_dims, const Kernel& kernel, const KernelArgsPackedArrayBase& args) { + const auto& cuda_kernel = static_cast(kernel); + ABSL_RETURN_IF_ERROR(cuda_kernel.UpdateMaxDynamicSharedMemoryBytes( + args.number_of_shared_bytes())); + + std::unique_ptr repacked; + const KernelArgsPackedArrayBase* packed_args = &args; + if (cuda_kernel.args_packing()) { + ABSL_ASSIGN_OR_RETURN(repacked, cuda_kernel.args_packing()(cuda_kernel, args)); + packed_args = repacked.get(); + } + + return CreateKernelNode( + dependencies, priority, threads, blocks, cluster_dims, + NativeKernel{cuda_kernel.gpu_function(), std::string(kernel.name()), + kernel.use_pdl()}, + *packed_args); +} + +absl::StatusOr CudaCommandBuffer::CreateKernelNode( + absl::Span dependencies, StreamPriority priority, + const ThreadDim& threads, const BlockDim& blocks, + const std::optional& cluster_dims, const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args) { const uint64_t shared_mem_bytes = args.number_of_shared_bytes(); XLA_VLOG_DEVICE(2, stream_exec_->device_ordinal()) @@ -733,19 +756,7 @@ absl::StatusOr CudaCommandBuffer::CreateKernelNode( << "): " << FormatGraphNodeHandles(dependencies); CUgraphNode node_handle = nullptr; - const auto& cuda_kernel = static_cast(kernel); - CUfunction function = cuda_kernel.gpu_function(); - ABSL_RETURN_IF_ERROR( - cuda_kernel.UpdateMaxDynamicSharedMemoryBytes(shared_mem_bytes)); - - std::unique_ptr repacked; - const KernelArgsPackedArrayBase* packed_args; - if (cuda_kernel.args_packing()) { - ABSL_ASSIGN_OR_RETURN(repacked, cuda_kernel.args_packing()(cuda_kernel, args)); - packed_args = repacked.get(); - } else { - packed_args = &args; - } + CUfunction function = static_cast(kernel.device_fn); auto set_params = [&](auto& params) { params.func = function; @@ -756,8 +767,10 @@ absl::StatusOr CudaCommandBuffer::CreateKernelNode( params.blockDimY = threads.y; params.blockDimZ = threads.z; params.sharedMemBytes = shared_mem_bytes; - params.kernelParams = - const_cast(packed_args->argument_addresses().data()); + // CUDA driver API requires void** for kernelParams even though it does not + // mutate the argument pointers. + // NOLINTNEXTLINE + params.kernelParams = const_cast(args.argument_addresses().data()); params.extra = nullptr; }; @@ -845,6 +858,28 @@ absl::Status CudaCommandBuffer::UpdateKernelNode( GraphNodeHandle node_handle, const ThreadDim& threads, const BlockDim& blocks, const std::optional& cluster_dims, const Kernel& kernel, const KernelArgsPackedArrayBase& args) { + const auto& cuda_kernel = static_cast(kernel); + ABSL_RETURN_IF_ERROR(cuda_kernel.UpdateMaxDynamicSharedMemoryBytes( + args.number_of_shared_bytes())); + + std::unique_ptr repacked; + const KernelArgsPackedArrayBase* packed_args = &args; + if (cuda_kernel.args_packing()) { + ABSL_ASSIGN_OR_RETURN(repacked, cuda_kernel.args_packing()(cuda_kernel, args)); + packed_args = repacked.get(); + } + + return UpdateKernelNode( + node_handle, threads, blocks, cluster_dims, + NativeKernel{cuda_kernel.gpu_function(), std::string(kernel.name()), + kernel.use_pdl()}, + *packed_args); +} + +absl::Status CudaCommandBuffer::UpdateKernelNode( + GraphNodeHandle node_handle, const ThreadDim& threads, + const BlockDim& blocks, const std::optional& cluster_dims, + const NativeKernel& kernel, const KernelArgsPackedArrayBase& args) { const uint64_t shared_mem_bytes = args.number_of_shared_bytes(); VLOG(2) << "Set kernel node params " << node_handle << " in graph executable " @@ -855,18 +890,7 @@ absl::Status CudaCommandBuffer::UpdateKernelNode( << "; shmem: " << shared_mem_bytes; CUDA_KERNEL_NODE_PARAMS params{}; - const auto& cuda_kernel = static_cast(kernel); - - std::unique_ptr repacked; - const KernelArgsPackedArrayBase* packed_args; - if (cuda_kernel.args_packing()) { - ABSL_ASSIGN_OR_RETURN(repacked, cuda_kernel.args_packing()(cuda_kernel, args)); - packed_args = repacked.get(); - } else { - packed_args = &args; - } - - CUfunction function = cuda_kernel.gpu_function(); + CUfunction function = static_cast(kernel.device_fn); params.func = function; params.gridDimX = blocks.x; params.gridDimY = blocks.y; @@ -875,13 +899,12 @@ absl::Status CudaCommandBuffer::UpdateKernelNode( params.blockDimY = threads.y; params.blockDimZ = threads.z; params.sharedMemBytes = shared_mem_bytes; - params.kernelParams = - const_cast(packed_args->argument_addresses().data()); + // CUDA driver API requires void** for kernelParams even though it does not + // mutate the argument pointers. + // NOLINTNEXTLINE + params.kernelParams = const_cast(args.argument_addresses().data()); params.extra = nullptr; - ABSL_RETURN_IF_ERROR( - cuda_kernel.UpdateMaxDynamicSharedMemoryBytes(shared_mem_bytes)); - if (cluster_dims.has_value()) { CUlaunchAttributeValue value; value.clusterDim.x = static_cast(cluster_dims->x); diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.h b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.h index 1145053f600854..48e4ced1630c75 100644 --- a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.h +++ b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.h @@ -173,6 +173,12 @@ class CudaCommandBuffer final : public GpuCommandBuffer { const std::optional& cluster_dims, const Kernel& kernel, const KernelArgsPackedArrayBase& args) override; + absl::StatusOr CreateKernelNode( + absl::Span dependencies, StreamPriority priority, + const ThreadDim& threads, const BlockDim& blocks, + const std::optional& cluster_dims, const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args) override; + absl::Status UpdateKernelNode(GraphNodeHandle node_handle, const ThreadDim& threads, const BlockDim& blocks, @@ -180,6 +186,13 @@ class CudaCommandBuffer final : public GpuCommandBuffer { const Kernel& kernel, const KernelArgsPackedArrayBase& args) override; + absl::Status UpdateKernelNode(GraphNodeHandle node_handle, + const ThreadDim& threads, + const BlockDim& blocks, + const std::optional& cluster_dims, + const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args) override; + absl::StatusOr CreateEmptyNode( absl::Span dependencies) override; diff --git a/third_party/xla/xla/stream_executor/gpu/BUILD b/third_party/xla/xla/stream_executor/gpu/BUILD index 463ae18dbb36be..6a4ee1dfccabe7 100644 --- a/third_party/xla/xla/stream_executor/gpu/BUILD +++ b/third_party/xla/xla/stream_executor/gpu/BUILD @@ -350,6 +350,31 @@ cc_library( "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/strings:string_view", + "@tsl//tsl/platform:path", + ], +) + +xla_test( + name = "gpu_hsaco_bundle_test", + srcs = ["gpu_hsaco_bundle_test.cc"], + backends = ["gpu"], + tags = ["rocm-only"], + deps = [ + ":asm_compiler", + "//xla:debug_options_flags", + "//xla/service/gpu/llvm_gpu_backend:amdgpu_backend", + "//xla/stream_executor:platform", + "//xla/stream_executor:platform_manager", + "//xla/stream_executor:stream_executor_h", + "//xla/stream_executor/rocm:rocm_platform", # buildcleaner: keep + "//xla/tsl/platform:env", + "//xla/tsl/platform:rocm_rocdl_path", + "//xla/tsl/platform:statusor", + "@com_google_googletest//:gtest_main", # buildcleaner: keep + "@llvm-project//llvm:AsmParser", + "@llvm-project//llvm:Support", + "@llvm-project//llvm:ir_headers", "@tsl//tsl/platform:path", ], ) diff --git a/third_party/xla/xla/stream_executor/gpu/asm_compiler.cc b/third_party/xla/xla/stream_executor/gpu/asm_compiler.cc index 70b955d9bbdd01..79a0cecf616723 100644 --- a/third_party/xla/xla/stream_executor/gpu/asm_compiler.cc +++ b/third_party/xla/xla/stream_executor/gpu/asm_compiler.cc @@ -58,13 +58,14 @@ absl::StatusOr> BundleGpuAsm( findRocmExecutable("llvm/bin/clang-offload-bundler", rocm_root_dir); // Initialise the "--inputs" / "--targets" arguments for the - // clang-offload-bundler with a dummy file / host target triple... - // clang-offload-bundler requires 1 and only 1 host target triple + // clang-offload-bundler with a dummy file / host target ID. + // clang-offload-bundler requires exactly one host target ID. + // Each ID is '----[-]'. std::ostringstream inputs_list; std::ostringstream targets_list; inputs_list << "/dev/null"; - targets_list << "host-x86_64-unknown-linux"; + targets_list << "host-x86_64-unknown-linux-gnu"; // Write images to temporary files. std::vector image_paths; @@ -79,7 +80,8 @@ absl::StatusOr> BundleGpuAsm( env, img_path, std::string(img.bytes.begin(), img.bytes.end()))); VLOG(2) << "image written to " << img_path; inputs_list << "," << img_path; - targets_list << ",hip-amdgcn-amd-amdhsa-" << img.gfx_arch; + // Empty is the extra '-' before gfx (hip-amdgcn-amd-amdhsa--gfx...). + targets_list << ",hip-amdgcn-amd-amdhsa--" << img.gfx_arch; image_paths.push_back(std::move(img_path)); } absl::Cleanup image_files_cleaner = [&image_paths] { diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.cc b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.cc index 0a06de913f30ea..e99b249e3c5f11 100644 --- a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.cc +++ b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.cc @@ -164,6 +164,31 @@ absl::Status GpuCommandBuffer::UpdateLaunchWithPackedArgs( kernel, packed_args); } +absl::StatusOr GpuCommandBuffer::CreateLaunch( + const ThreadDim& threads, const BlockDim& blocks, + const std::optional& cluster_dims, const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args, + absl::Span dependencies, StreamPriority priority) { + ABSL_RETURN_IF_ERROR(CheckInState(State::kCreate)); + + ABSL_ASSIGN_OR_RETURN( + GraphNodeHandle handle, + CreateKernelNode(ToGraphNodeDependencies(dependencies), priority, threads, + blocks, cluster_dims, kernel, args)); + + return AppendCommand(GpuCommand{handle}); +} + +absl::Status GpuCommandBuffer::UpdateLaunch( + const Command* command, const ThreadDim& threads, const BlockDim& blocks, + const std::optional& cluster_dims, const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args) { + ABSL_RETURN_IF_ERROR(CheckInState(State::kUpdate)); + auto* gpu_command = absl::down_cast(command); + return UpdateKernelNode(gpu_command->handle, threads, blocks, cluster_dims, + kernel, args); +} + absl::StatusOr GpuCommandBuffer::CreateLaunch( const ThreadDim& threads, const BlockDim& blocks, const std::optional& cluster_dims, const Kernel& kernel, diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.h b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.h index 3053c21b5dcb0f..bcb43b5272c230 100644 --- a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.h +++ b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.h @@ -135,12 +135,25 @@ class GpuCommandBuffer : public CommandBuffer { const KernelArgs& args, absl::Span dependencies, StreamPriority priority) override; + absl::StatusOr CreateLaunch( + const ThreadDim& threads, const BlockDim& blocks, + const std::optional& cluster_dims, const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args, + absl::Span dependencies, + StreamPriority priority) override; + absl::Status UpdateLaunch(const Command* command, const ThreadDim& threads, const BlockDim& blocks, const std::optional& cluster_dims, const Kernel& kernel, const KernelArgs& args) override; + absl::Status UpdateLaunch(const Command* command, const ThreadDim& threads, + const BlockDim& blocks, + const std::optional& cluster_dims, + const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args) override; + absl::StatusOr CreateChildCommand( const CommandBuffer& nested, absl::Span dependencies) override; @@ -420,6 +433,12 @@ class GpuCommandBuffer : public CommandBuffer { const std::optional& cluster_dims, const Kernel& kernel, const KernelArgsPackedArrayBase& args) = 0; + virtual absl::StatusOr CreateKernelNode( + absl::Span dependencies, StreamPriority priority, + const ThreadDim& threads, const BlockDim& blocks, + const std::optional& cluster_dims, const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args) = 0; + // Updates the kernel launch node with the given parameters. Will return an // error if the given node has not been created as a kernel launch node. virtual absl::Status UpdateKernelNode( @@ -427,6 +446,11 @@ class GpuCommandBuffer : public CommandBuffer { const BlockDim& blocks, const std::optional& cluster_dims, const Kernel& kernel, const KernelArgsPackedArrayBase& args) = 0; + virtual absl::Status UpdateKernelNode( + GraphNodeHandle node_handle, const ThreadDim& threads, + const BlockDim& blocks, const std::optional& cluster_dims, + const NativeKernel& kernel, const KernelArgsPackedArrayBase& args) = 0; + virtual absl::StatusOr CreateHostNode( absl::Span dependencies, absl::AnyInvocable callback) = 0; diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_hsaco_bundle_test.cc b/third_party/xla/xla/stream_executor/gpu/gpu_hsaco_bundle_test.cc new file mode 100644 index 00000000000000..cb5804eda665eb --- /dev/null +++ b/third_party/xla/xla/stream_executor/gpu/gpu_hsaco_bundle_test.cc @@ -0,0 +1,106 @@ +/* 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 +#include +#include + +#include +#include +#include "llvm/AsmParser/Parser.h" +#include "llvm/IR/LLVMContext.h" +#include "llvm/IR/Module.h" +#include "llvm/Support/SourceMgr.h" +#include "xla/debug_options_flags.h" +#include "xla/service/gpu/llvm_gpu_backend/amdgpu_backend.h" +#include "xla/stream_executor/gpu/asm_compiler.h" +#include "xla/stream_executor/platform.h" +#include "xla/stream_executor/platform_manager.h" +#include "xla/stream_executor/stream_executor.h" +#include "xla/tsl/platform/env.h" +#include "xla/tsl/platform/rocm_rocdl_path.h" +#include "xla/tsl/platform/statusor.h" +#include "tsl/platform/path.h" + +namespace stream_executor { +namespace { + +TEST(GpuHsacoBundleTest, CompileToHsacoThenBundleGpuAsmWorks) { + const std::string rocm_root = tsl::RocmRoot(); + if (rocm_root.empty()) { + GTEST_SKIP() << "ROCm root is empty; clang-offload-bundler is unavailable"; + } + const std::string bundler = + tsl::io::JoinPath(rocm_root, "llvm/bin/clang-offload-bundler"); + if (!tsl::Env::Default()->FileExists(bundler).ok()) { + GTEST_SKIP() << "clang-offload-bundler not found at " << bundler; + } + + ASSERT_OK_AND_ASSIGN(Platform * platform, + PlatformManager::PlatformWithName("ROCM")); + ASSERT_GT(platform->VisibleDeviceCount(), 0); + ASSERT_OK_AND_ASSIGN(StreamExecutor * executor, + platform->ExecutorForDevice(0)); + const std::string gfx_arch = + executor->GetDeviceDescription().rocm_compute_capability().gfx_version(); + ASSERT_FALSE(gfx_arch.empty()); + + // CompileToHsaco takes an LLVM module; this kernel has no host or ROCDL deps. + constexpr char kIr[] = R"( +target triple = "amdgcn-amd-amdhsa" +define amdgpu_kernel void @simple_add(ptr addrspace(1) %a, ptr addrspace(1) %b, + ptr addrspace(1) %c) { +entry: + %tid = call i32 @llvm.amdgcn.workitem.id.x() + %tidx = zext i32 %tid to i64 + %a_ptr = getelementptr float, ptr addrspace(1) %a, i64 %tidx + %b_ptr = getelementptr float, ptr addrspace(1) %b, i64 %tidx + %c_ptr = getelementptr float, ptr addrspace(1) %c, i64 %tidx + %a_val = load float, ptr addrspace(1) %a_ptr, align 4 + %b_val = load float, ptr addrspace(1) %b_ptr, align 4 + %sum = fadd float %a_val, %b_val + store float %sum, ptr addrspace(1) %c_ptr, align 4 + ret void +} +declare i32 @llvm.amdgcn.workitem.id.x() +)"; + + llvm::LLVMContext llvm_context; + llvm::SMDiagnostic parse_err; + std::unique_ptr module = + llvm::parseAssemblyString(kIr, parse_err, llvm_context); + ASSERT_NE(module, nullptr) << parse_err.getMessage().str(); + // EmitModuleToHsaco joins the module id into a temp path. + module->setModuleIdentifier("gpu_hsaco_bundle_test"); + + ASSERT_OK_AND_ASSIGN( + xla::gpu::amdgpu::HsacoResult hsaco_result, + xla::gpu::amdgpu::CompileToHsaco( + module.get(), + executor->GetDeviceDescription().gpu_compute_capability(), + xla::GetDebugOptionsFromFlags(), "gpu_hsaco_bundle_test")); + ASSERT_FALSE(hsaco_result.hsaco.empty()); + + HsacoImage image; + image.gfx_arch = gfx_arch; + image.bytes = std::move(hsaco_result.hsaco); + ASSERT_OK_AND_ASSIGN(std::vector bundle, + BundleGpuAsm({image}, rocm_root)); + ASSERT_FALSE(bundle.empty()); +} + +} // namespace +} // namespace stream_executor diff --git a/third_party/xla/xla/stream_executor/kernel.h b/third_party/xla/xla/stream_executor/kernel.h index a754b67e243d2f..77797c10c4d042 100644 --- a/third_party/xla/xla/stream_executor/kernel.h +++ b/third_party/xla/xla/stream_executor/kernel.h @@ -92,6 +92,17 @@ namespace stream_executor { // Kernel //===----------------------------------------------------------------------===// +// Platform-agnostic wrapper for an externally-managed native device function +// pointer (e.g. CUfunction or hipFunction_t). +struct NativeKernel { + void* device_fn = nullptr; + std::string kernel_name; + bool pdl = false; + + absl::string_view name() const { return kernel_name; } + bool use_pdl() const { return pdl; } +}; + // A data-parallel kernel (code entity) for launching via the StreamExecutor, // analogous to a void* device function pointer. See TypedKernel for the typed // variant. diff --git a/third_party/xla/xla/stream_executor/mock_command_buffer.h b/third_party/xla/xla/stream_executor/mock_command_buffer.h index 15aae998c0c2c2..1f830c44e9fbb8 100644 --- a/third_party/xla/xla/stream_executor/mock_command_buffer.h +++ b/third_party/xla/xla/stream_executor/mock_command_buffer.h @@ -55,12 +55,27 @@ class MockCommandBuffer : public CommandBuffer { absl::Span dependencies, StreamPriority priority), (override)); + MOCK_METHOD(absl::StatusOr, CreateLaunch, + (const ThreadDim& threads, const BlockDim& blocks, + const std::optional& cluster_dims, + const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args, + absl::Span dependencies, + StreamPriority priority), + (override)); MOCK_METHOD(absl::Status, UpdateLaunch, (const Command* command, const ThreadDim& threads, const BlockDim& blocks, const std::optional& cluster_dims, const Kernel& kernel, const KernelArgs& args), (override)); + MOCK_METHOD(absl::Status, UpdateLaunch, + (const Command* command, const ThreadDim& threads, + const BlockDim& blocks, + const std::optional& cluster_dims, + const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args), + (override)); MOCK_METHOD(absl::StatusOr, CreateChildCommand, (const CommandBuffer& nested, absl::Span dependencies), diff --git a/third_party/xla/xla/stream_executor/rocm/BUILD b/third_party/xla/xla/stream_executor/rocm/BUILD index 0a73c4da4aa0cc..9a87e4ba0e043b 100644 --- a/third_party/xla/xla/stream_executor/rocm/BUILD +++ b/third_party/xla/xla/stream_executor/rocm/BUILD @@ -1233,6 +1233,7 @@ cc_library( "//xla/stream_executor:command_buffer", "//xla/stream_executor:device_address", "//xla/stream_executor:kernel", + "//xla/stream_executor:kernel_args", "//xla/stream_executor:launch_dim", "//xla/stream_executor:platform", "//xla/stream_executor:stream_executor_h", diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.cc b/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.cc index b81487137e927d..3e8d3aef633e9e 100644 --- a/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.cc +++ b/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.cc @@ -39,6 +39,7 @@ limitations under the License. #include "xla/stream_executor/device_address.h" #include "xla/stream_executor/gpu/gpu_command_buffer.h" #include "xla/stream_executor/kernel.h" +#include "xla/stream_executor/kernel_args.h" #include "xla/stream_executor/launch_dim.h" #include "xla/stream_executor/platform.h" #include "xla/stream_executor/rocm/rocm_kernel.h" @@ -327,6 +328,25 @@ absl::StatusOr RocmCommandBuffer::CreateKernelNode( const ThreadDim& threads, const BlockDim& blocks, const std::optional& cluster_dims, const Kernel& kernel, const KernelArgsPackedArrayBase& args) { + std::unique_ptr repacked; + const KernelArgsPackedArrayBase* packed_args = &args; + if (kernel.args_packing()) { + ABSL_ASSIGN_OR_RETURN(repacked, kernel.args_packing()(kernel, args)); + packed_args = repacked.get(); + } + + return CreateKernelNode( + dependencies, priority, threads, blocks, cluster_dims, + NativeKernel{static_cast(kernel).gpu_function(), + std::string(kernel.name()), kernel.use_pdl()}, + *packed_args); +} + +absl::StatusOr RocmCommandBuffer::CreateKernelNode( + absl::Span dependencies, StreamPriority priority, + const ThreadDim& threads, const BlockDim& blocks, + const std::optional& cluster_dims, const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args) { const uint64_t shared_mem_bytes = args.number_of_shared_bytes(); VLOG(2) << "Add kernel node to a graph " << graph_ @@ -336,18 +356,8 @@ absl::StatusOr RocmCommandBuffer::CreateKernelNode( << " bdz: " << threads.z << "; shmem: " << shared_mem_bytes << "; deps: " << dependencies.size(); - std::unique_ptr repacked; - const KernelArgsPackedArrayBase* packed_args; - if (kernel.args_packing()) { - ABSL_ASSIGN_OR_RETURN(repacked, kernel.args_packing()(kernel, args)); - packed_args = repacked.get(); - } else { - packed_args = &args; - } - hipKernelNodeParams params{}; - hipFunction_t function = - static_cast(kernel).gpu_function(); + hipFunction_t function = static_cast(kernel.device_fn); params.func = function; params.gridDim.x = blocks.x; params.gridDim.y = blocks.y; @@ -356,8 +366,10 @@ absl::StatusOr RocmCommandBuffer::CreateKernelNode( params.blockDim.y = threads.y; params.blockDim.z = threads.z; params.sharedMemBytes = shared_mem_bytes; - params.kernelParams = - const_cast(packed_args->argument_addresses().data()); + // HIP driver API requires void** for kernelParams even though it does not + // mutate the argument pointers. + // NOLINTNEXTLINE + params.kernelParams = const_cast(args.argument_addresses().data()); params.extra = nullptr; if (shared_mem_bytes != 0) { @@ -383,6 +395,24 @@ absl::Status RocmCommandBuffer::UpdateKernelNode( GraphNodeHandle node_handle, const ThreadDim& threads, const BlockDim& blocks, const std::optional& cluster_dims, const Kernel& kernel, const KernelArgsPackedArrayBase& args) { + std::unique_ptr repacked; + const KernelArgsPackedArrayBase* packed_args = &args; + if (kernel.args_packing()) { + ABSL_ASSIGN_OR_RETURN(repacked, kernel.args_packing()(kernel, args)); + packed_args = repacked.get(); + } + + return UpdateKernelNode( + node_handle, threads, blocks, cluster_dims, + NativeKernel{static_cast(kernel).gpu_function(), + std::string(kernel.name()), kernel.use_pdl()}, + *packed_args); +} + +absl::Status RocmCommandBuffer::UpdateKernelNode( + GraphNodeHandle node_handle, const ThreadDim& threads, + const BlockDim& blocks, const std::optional& cluster_dims, + const NativeKernel& kernel, const KernelArgsPackedArrayBase& args) { const uint64_t shared_mem_bytes = args.number_of_shared_bytes(); VLOG(2) << "Set kernel node params " << node_handle << " in graph executable " @@ -391,18 +421,8 @@ absl::Status RocmCommandBuffer::UpdateKernelNode( << " bdx: " << threads.x << " bdy: " << threads.y << " bdz: " << threads.z << "; shmem: " << shared_mem_bytes; - std::unique_ptr repacked; - const KernelArgsPackedArrayBase* packed_args; - if (kernel.args_packing()) { - ABSL_ASSIGN_OR_RETURN(repacked, kernel.args_packing()(kernel, args)); - packed_args = repacked.get(); - } else { - packed_args = &args; - } - hipKernelNodeParams params{}; - hipFunction_t function = - static_cast(kernel).gpu_function(); + hipFunction_t function = static_cast(kernel.device_fn); params.func = function; params.gridDim.x = blocks.x; params.gridDim.y = blocks.y; @@ -411,8 +431,10 @@ absl::Status RocmCommandBuffer::UpdateKernelNode( params.blockDim.y = threads.y; params.blockDim.z = threads.z; params.sharedMemBytes = shared_mem_bytes; - params.kernelParams = - const_cast(packed_args->argument_addresses().data()); + // HIP driver API requires void** for kernelParams even though it does not + // mutate the argument pointers. + // NOLINTNEXTLINE + params.kernelParams = const_cast(args.argument_addresses().data()); params.extra = nullptr; if (shared_mem_bytes != 0) { diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.h b/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.h index b8e3ac54e6016a..3e2fd7b35d145c 100644 --- a/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.h +++ b/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.h @@ -159,6 +159,12 @@ class RocmCommandBuffer : public GpuCommandBuffer { const std::optional& cluster_dims, const Kernel& kernel, const KernelArgsPackedArrayBase& args) override; + absl::StatusOr CreateKernelNode( + absl::Span dependencies, StreamPriority priority, + const ThreadDim& threads, const BlockDim& blocks, + const std::optional& cluster_dims, const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args) override; + absl::Status UpdateKernelNode(GraphNodeHandle node_handle, const ThreadDim& threads, const BlockDim& blocks, @@ -166,6 +172,13 @@ class RocmCommandBuffer : public GpuCommandBuffer { const Kernel& kernel, const KernelArgsPackedArrayBase& args) override; + absl::Status UpdateKernelNode(GraphNodeHandle node_handle, + const ThreadDim& threads, + const BlockDim& blocks, + const std::optional& cluster_dims, + const NativeKernel& kernel, + const KernelArgsPackedArrayBase& args) override; + absl::StatusOr CreateEmptyNode( absl::Span dependencies) override; diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/BUILD b/third_party/xla/xla/tests/aot_compatibility_experimental/BUILD index fa8e84382ea258..b35d17d3e5fcbf 100644 --- a/third_party/xla/xla/tests/aot_compatibility_experimental/BUILD +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/BUILD @@ -15,12 +15,17 @@ cc_library( deps = [ "//xla:xla_proto_cc", "//xla/pjrt:pjrt_client", + "//xla/stream_executor:platform", + "//xla/stream_executor:platform_manager", + "//xla/stream_executor:stream_executor_h", "//xla/tests:aot_interception_pjrt_client", "//xla/tests:hlo_test_base", "//xla/tests:pjrt_client_registry", "//xla/tsl/platform:env", "//xla/tsl/platform:logging", "//xla/tsl/platform:test", + "@com_google_absl//absl/log", + "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/BUILD b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/BUILD index 944a9784e441ba..6b59ed77da4a0f 100644 --- a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/BUILD +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/BUILD @@ -24,7 +24,7 @@ xla_test( name = "collective_ops_aot_test", srcs = ["collective_ops_aot_test.cc"], backends = _GPU_TARGETS, - data = glob(["executables/collective_ops_aot_test_2gpu/**/*.pbtxt"]), + data = glob(["executables/collective_ops_aot_test_2gpu_*/**/*.pbtxt"]), disabled_backends = [ # TODO(b/491194726): Serialized executable compute capability 6.0 does not match sm_10.3a. "gb300", diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/collective_ops_aot_test.cc b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/collective_ops_aot_test.cc index 6dcdf613e6804e..95f19a45dbee22 100644 --- a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/collective_ops_aot_test.cc +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/collective_ops_aot_test.cc @@ -142,7 +142,8 @@ std::vector GetTestParamsOrDie( INSTANTIATE_TEST_SUITE_P( BackwardsCompatibility, CollectiveOpsAotTest, ValuesIn(GetTestParamsOrDie(GetAotTestParamsForBackwardsCompatibility( - "collective_ops_aot_test_2gpu", AOTTestPlatform::kGpu))), + absl::StrCat("collective_ops_aot_test_2gpu_", DetectGpuArchToken()), + AOTTestPlatform::kGpu))), [](const TestParamInfo& info) { return absl::StrCat("v", info.param.version); }); @@ -150,7 +151,8 @@ INSTANTIATE_TEST_SUITE_P( INSTANTIATE_TEST_SUITE_P( GoldenFileVerification, CollectiveOpsAotTest, ValuesIn(GetTestParamsOrDie(GetAotTestParamsForGoldenFileVerification( - "collective_ops_aot_test_2gpu", AOTTestPlatform::kGpu))), + absl::StrCat("collective_ops_aot_test_2gpu_", DetectGpuArchToken()), + AOTTestPlatform::kGpu))), [](const TestParamInfo& info) { return absl::StrCat("v", info.param.version); }); diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v1/AllGatherMixedTypes.pbtxt b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v1/AllGatherMixedTypes.pbtxt new file mode 100644 index 00000000000000..4e7a607fc3ae22 --- /dev/null +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v1/AllGatherMixedTypes.pbtxt @@ -0,0 +1,2180 @@ +# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto +# proto-message: xla.HumanReadableAotExecutable +executable_and_options { + compile_options { + executable_build_options { + device_ordinal: -1 + result_layout { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.01 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: false + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: false + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + } + num_replicas: 2 + num_partitions: 1 + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + comp_envs { + } + allow_spmd_sharding_propagation_to_parameters: false + process_count: 1 + memory_fitting_level: EFFORT_O2 + } + } + pjrt_client_name: "PjRtStreamExecutorClient" +} +gpu_executable { + hlo_module_with_config { + hlo_module { + name: "test" + entry_computation_name: "test_computation" + computations { + name: "fused_transpose" + instructions { + name: "param_0.1" + opcode: "parameter" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_0.1" + } + id: 4294967298 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.6.1" + opcode: "bitcast" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.6.1" + } + id: 4294967297 + operand_ids: 4294967298 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "transpose.1.1" + opcode: "transpose" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "transpose.1.1" + } + dimensions: 1 + dimensions: 0 + id: 4294967296 + operand_ids: 4294967297 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + result { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.1" + } + id: 1 + root_id: 4294967296 + is_fusion_computation: true + } + computations { + name: "fused_transpose.1" + instructions { + name: "param_0.3" + opcode: "parameter" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_0.3" + } + id: 8589934594 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.4.1" + opcode: "bitcast" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.4.1" + } + id: 8589934593 + operand_ids: 8589934594 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "transpose.3" + opcode: "transpose" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "transpose.3" + } + dimensions: 1 + dimensions: 0 + id: 8589934592 + operand_ids: 8589934593 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + result { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.3" + } + id: 2 + root_id: 8589934592 + is_fusion_computation: true + } + computations { + name: "fused_broadcast" + instructions { + name: "param_0.5" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "param_0.5" + } + id: 12884901890 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "p1.1.1" + opcode: "convert" + shape { + element_type: F32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "p1.1.1" + } + id: 12884901889 + operand_ids: 12884901890 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "broadcast.13.1" + opcode: "broadcast" + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "broadcast.13.1" + } + id: 12884901888 + operand_ids: 12884901889 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.5" + } + id: 3 + root_id: 12884901888 + is_fusion_computation: true + } + computations { + name: "wrapped_broadcast_computation" + instructions { + name: "param_0.7" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "param_0.7" + } + id: 21474836481 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "broadcast.12.1" + opcode: "broadcast" + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "broadcast.12.1" + } + id: 21474836480 + operand_ids: 21474836481 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.7" + } + id: 5 + root_id: 21474836480 + is_fusion_computation: true + } + computations { + name: "test_computation" + instructions { + name: "id.1" + opcode: "replica-id" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "id.1" + } + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "wrapped_broadcast" + opcode: "fusion" + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "wrapped_broadcast" + } + fusion_kind: "kLoop" + id: 9 + operand_ids: 0 + called_computation_ids: 5 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_broadcast_fusion" + opcode: "fusion" + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_broadcast_fusion" + } + fusion_kind: "kLoop" + id: 8 + operand_ids: 0 + called_computation_ids: 3 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "allgather.3" + opcode: "all-gather" + shape { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + metadata { + scheduling_name: "all-gather-start" + } + dimensions: 1 + id: 10 + operand_ids: 9 + operand_ids: 8 + frontend_attributes { + } + statistics_viz { + } + collective_device_list { + } + backend_config_payload { + id: 0 + } + } + instructions { + name: "ag0.1" + opcode: "get-tuple-element" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "ag0.1" + } + id: 1 + operand_ids: 10 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_transpose_fusion.1" + opcode: "fusion" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_transpose_fusion.1" + } + fusion_kind: "kLoop" + id: 7 + operand_ids: 1 + called_computation_ids: 2 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.24" + opcode: "bitcast" + shape { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.24" + } + id: 4 + operand_ids: 7 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "ag1.1" + opcode: "get-tuple-element" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "ag1.1" + } + tuple_index: 1 + id: 2 + operand_ids: 10 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_transpose_fusion" + opcode: "fusion" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_transpose_fusion" + } + fusion_kind: "kLoop" + id: 6 + operand_ids: 2 + called_computation_ids: 1 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.1.0" + opcode: "bitcast" + shape { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.1.0" + } + id: 5 + operand_ids: 6 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "out.1" + opcode: "tuple" + shape { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + metadata { + scheduling_name: "out.1" + } + id: 3 + operand_ids: 4 + operand_ids: 5 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + root_id: 3 + } + host_program_shape { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + schedule { + sequences { + key: 0 + value { + instruction_ids: 0 + instruction_ids: 8 + instruction_ids: 9 + instruction_ids: 10 + instruction_ids: 1 + instruction_ids: 7 + instruction_ids: 2 + instruction_ids: 6 + instruction_ids: 4 + instruction_ids: 5 + instruction_ids: 3 + } + } + } + input_output_alias { + } + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + buffer_donor { + } + frontend_attributes { + map { + key: "fingerprint_before_lhs" + value: "7b2e4f79c5524b5cbc566da178aa7fda" + } + } + payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" + } + config { + entry_computation_layout { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + replica_count: 2 + num_partitions: 1 + intra_op_parallelism_threads: -1 + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.01 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: false + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: false + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + } + static_device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + allow_spmd_sharding_propagation_to_parameters: false + sharding_config { + } + memory_fitting_level: EFFORT_O2 + schedule_config { + } + partition_size: 2 + } + } + buffer_assignment { + logical_buffers { + id: 9 + size: 4 + defined_at { + } + } + logical_buffers { + id: 10 + size: 8 + defined_at { + instruction_id: 9 + } + } + logical_buffers { + id: 11 + size: 8 + defined_at { + instruction_id: 8 + } + } + logical_buffers { + id: 12 + size: 16 + defined_at { + instruction_id: 10 + } + } + logical_buffers { + id: 13 + size: 16 + defined_at { + shape_index: 0 + instruction_id: 10 + } + } + logical_buffers { + id: 14 + size: 16 + defined_at { + shape_index: 1 + instruction_id: 10 + } + } + logical_buffers { + id: 15 + size: 16 + defined_at { + instruction_id: 7 + } + } + logical_buffers { + id: 16 + size: 16 + defined_at { + instruction_id: 6 + } + } + logical_buffers { + id: 17 + size: 16 + defined_at { + instruction_id: 3 + } + } + buffer_allocations { + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 10 + size: 8 + element_type: U32 + } + assigned { + logical_buffer_id: 15 + size: 16 + element_type: U32 + } + } + buffer_allocations { + index: 1 + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 9 + size: 4 + element_type: U32 + } + assigned { + logical_buffer_id: 13 + size: 16 + element_type: U32 + } + assigned { + logical_buffer_id: 16 + size: 16 + element_type: F32 + } + } + buffer_allocations { + index: 2 + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 17 + size: 16 + element_type: TUPLE + } + is_tuple: true + } + buffer_allocations { + index: 3 + size: 520 + assigned { + logical_buffer_id: 11 + offset: 512 + size: 8 + element_type: F32 + } + assigned { + logical_buffer_id: 12 + size: 16 + element_type: TUPLE + } + assigned { + logical_buffer_id: 14 + offset: 256 + size: 16 + element_type: F32 + } + is_tuple: true + } + heap_simulator_traces { + events { + buffer_id: 11 + computation_name: "test_computation" + instruction_name: "loop_broadcast_fusion" + } + events { + buffer_id: 14 + computation_name: "test_computation" + instruction_name: "allgather.3" + } + events { + kind: FREE + buffer_id: 11 + computation_name: "test_computation" + instruction_name: "loop_broadcast_fusion" + } + events { + kind: FREE + buffer_id: 14 + computation_name: "test_computation" + instruction_name: "allgather.3" + } + whole_module_simulation: true + buffer_allocation_index: 3 + } + peak_buffers { + logical_buffer_ids: 15 + } + peak_buffers { + index: 1 + logical_buffer_ids: 16 + } + peak_buffers { + index: 2 + logical_buffer_ids: 17 + } + peak_buffers { + index: 3 + logical_buffer_ids: 12 + logical_buffer_ids: 11 + logical_buffer_ids: 14 + } + } + gpu_compute_capability { + cuda_compute_capability { + major: 10 + feature_extension: ACCELERATED_FEATURES + } + } + module_name: "test" + program_shape { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + buffer_allocations { + values { + size: 16 + maybe_live_out: true + } + values { + index: 1 + size: 16 + maybe_live_out: true + } + values { + index: 2 + size: 16 + maybe_live_out: true + is_tuple: true + } + values { + index: 3 + size: 520 + is_tuple: true + } + } + output_info_map { + shape_index { + } + output_info { + allocation_index: 2 + } + } + output_info_map { + shape_index { + indexes: 0 + } + output_info { + } + } + output_info_map { + shape_index { + indexes: 1 + } + output_info { + allocation_index: 1 + } + } + thunks { + thunk_info { + profile_annotation: "id.1" + thunk_id: 1 + } + replica_id_thunk { + dest_buffer { + size: 4 + buffer_allocation_index: 1 + element_type: U32 + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_broadcast_fusion" + thunk_id: 2 + } + custom_kernel_thunk { + args { + slice { + size: 4 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + } + args { + slice { + offset: 512 + size: 8 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "loop_broadcast_fusion" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_broadcast_fusion\000.nv.capmerc.text.loop_broadcast_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_broadcast_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_broadcast_fusion\000.nv.constant0.loop_broadcast_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000j\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004(\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\240\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\340\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000\"\016\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\342\000\000\000\n\000\000\000b\016\000\201y\002\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\342/\000Er\007\000\002\000\000\000\000\020 \000\000\312O\000\206y\000\004\007\000\000\000\004\031\020\014\000\342\017\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\014\000\000@\300\005\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\001\000\000\000\000\000\000\000\000\000\000\000\320\006\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\004\020\000\000\000\004\360\000\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\360\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\372\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000^\002\000\000\000\000\000\000\216\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\316\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\030\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\314\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000m\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\006\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\303\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\376\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\354\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\032\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000e\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000z\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\213\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\262\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\314\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\352\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "loop_broadcast_fusion" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 2 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "wrapped_broadcast" + thunk_id: 3 + } + custom_kernel_thunk { + args { + slice { + size: 4 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + } + args { + slice { + size: 8 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "wrapped_broadcast" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.wrapped_broadcast\000.nv.capmerc.text.wrapped_broadcast\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.wrapped_broadcast\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000wrapped_broadcast\000.nv.constant0.wrapped_broadcast\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\320\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\374\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000Z\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004$\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\220\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\340\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000\"\016\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\342\000\000\000\n\000\000\000b\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\312/\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\013\000\000@\300\002\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\007\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\000\000\000\000\000\000\000\004\020\000\000\000\004\340\000\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\340\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\320\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\374\000\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\360\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\342\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000F\002\000\000\000\000\000\000z\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\302\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\350\004\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\234\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000i\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\344\005\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\267\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\362\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000X\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\340\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000f\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000w\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\232\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\264\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\322\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "wrapped_broadcast" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 2 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "allgather.3" + thunk_id: 4 + } + all_gather_thunk { + collective_config { + operand_element_type: U32 + operand_element_type: F32 + } + buffers { + element_count: 2 + source_buffer { + slice { + size: 8 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + destination_buffer { + slice { + size: 16 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + } + buffers { + element_count: 2 + source_buffer { + slice { + offset: 512 + size: 8 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + destination_buffer { + slice { + offset: 256 + size: 16 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_transpose_fusion.1" + thunk_id: 5 + } + custom_kernel_thunk { + args { + slice { + size: 16 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + args { + slice { + size: 16 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "loop_transpose_fusion_1" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion_1\000.nv.capmerc.text.loop_transpose_fusion_1\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_transpose_fusion_1\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion_1\000.nv.constant0.loop_transpose_fusion_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\342\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Z\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000r\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\340\000\000\000\n\000\000\000b\016\000\254w\004\377\000k\000\000\000\n\000\010\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\342\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242N\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\017\000\000@\340-\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\005\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\004\020\000\000\000\004`\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000`\001\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\342\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Z\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\006\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000j\002\000\000\000\000\000\000\230\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\004\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\324\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\0000\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\344\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\006\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000o\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\006\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\311\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\001\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\362\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000F\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000o\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\204\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\225\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\276\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\330\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "loop_transpose_fusion_1" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 4 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_transpose_fusion" + thunk_id: 6 + } + custom_kernel_thunk { + args { + slice { + offset: 256 + size: 16 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + args { + slice { + size: 16 + buffer_allocation_index: 1 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "loop_transpose_fusion" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion\000.nv.capmerc.text.loop_transpose_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_transpose_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion\000.nv.constant0.loop_transpose_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000j\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\340\000\000\000\n\000\000\000b\016\000\254w\004\377\000k\000\000\000\n\000\010\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\342\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242N\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\017\000\000@\340-\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\005\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\004\020\000\000\000\004`\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000`\001\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\372\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000^\002\000\000\000\000\000\000\216\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\316\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\030\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\314\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000m\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\006\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\303\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\376\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\354\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\032\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000e\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000z\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\213\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\262\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\314\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\352\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "loop_transpose_fusion" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 4 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + executable_abi_version { + platform_name: "CUDA" + cuda_platform_version { + cuda_toolkit_version: "12.8.0" + cudnn_version: "9.10.2" + cub_version: "3.1.2" + } + } + buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 0B\n maybe_live_out allocation: 48B\n preallocated temp allocation: 520B\n preallocated temp fragmentation: 240B (46.15%)\n total allocation: 568B\n total fragmentation: 504B (88.73%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: (u32[2,2], f32[2,2])\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 16B\n\t\tXLA Label: tuple\n\t\tShape: (u32[4], f32[4])\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: u32[2,2]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 8B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,1]\n\t\t==========================\n" +} diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v1/ReduceScatter.pbtxt b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v1/ReduceScatter.pbtxt new file mode 100644 index 00000000000000..0f7f9346f06368 --- /dev/null +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v1/ReduceScatter.pbtxt @@ -0,0 +1,1615 @@ +# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto +# proto-message: xla.HumanReadableAotExecutable +executable_and_options { + compile_options { + executable_build_options { + device_ordinal: -1 + result_layout { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.01 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: false + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: false + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + } + num_replicas: 2 + num_partitions: 1 + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + comp_envs { + } + allow_spmd_sharding_propagation_to_parameters: false + process_count: 1 + memory_fitting_level: EFFORT_O2 + } + } + pjrt_client_name: "PjRtStreamExecutorClient" +} +gpu_executable { + hlo_module_with_config { + hlo_module { + name: "test" + entry_computation_name: "main" + computations { + name: "add" + instructions { + name: "lhs" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "lhs" + } + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "rhs" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "rhs" + } + parameter_number: 1 + id: 1 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "add.1" + opcode: "add" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "add.1" + } + id: 2 + operand_ids: 0 + operand_ids: 1 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + parameter_names: "lhs" + parameter_names: "rhs" + } + root_id: 2 + } + computations { + name: "fused_select" + instructions { + name: "param_2.2" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "param_2.2" + } + parameter_number: 2 + id: 8589934597 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "zero_2" + opcode: "constant" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + } + literal { + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + u32s: 0 + } + id: 8589934598 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "p.2" + opcode: "compare" + shape { + element_type: PRED + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "p.2" + } + id: 8589934596 + operand_ids: 8589934597 + operand_ids: 8589934598 + comparison_direction: "EQ" + frontend_attributes { + } + comparison_type: "UNSIGNED" + statistics_viz { + } + comparison_order: "TOTAL" + } + instructions { + name: "pb.2" + opcode: "broadcast" + shape { + element_type: PRED + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "pb.2" + } + id: 8589934593 + operand_ids: 8589934596 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "param_0.1" + opcode: "parameter" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_0.1" + } + id: 8589934594 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "param_1.1" + opcode: "parameter" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_1.1" + } + parameter_number: 1 + id: 8589934595 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "data.2" + opcode: "select" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "data.2" + } + id: 8589934592 + operand_ids: 8589934593 + operand_ids: 8589934594 + operand_ids: 8589934595 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + parameters { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + parameter_names: "param_0.1" + parameter_names: "param_1.1" + parameter_names: "param_2.2" + } + id: 2 + root_id: 8589934592 + is_fusion_computation: true + } + computations { + name: "main" + instructions { + name: "c0_1" + opcode: "constant" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + } + literal { + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + u32s: 1 + u32s: 2 + u32s: 3 + u32s: 4 + u32s: 5 + u32s: 6 + u32s: 7 + u32s: 8 + } + id: 4294967297 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "c1_1" + opcode: "constant" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + } + literal { + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + u32s: 10 + u32s: 11 + u32s: 12 + u32s: 13 + u32s: 14 + u32s: 15 + u32s: 16 + u32s: 17 + } + id: 4294967298 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "id.1" + opcode: "replica-id" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "id.1" + } + id: 4294967296 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_select_fusion" + opcode: "fusion" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_select_fusion" + } + fusion_kind: "kLoop" + id: 4294967299 + operand_ids: 4294967297 + operand_ids: 4294967298 + operand_ids: 4294967296 + called_computation_ids: 2 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "ars.3" + opcode: "reduce-scatter" + shape { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "reduce-scatter-start" + } + dimensions: 0 + id: 4294967300 + operand_ids: 4294967299 + called_computation_ids: 0 + frontend_attributes { + } + statistics_viz { + } + collective_device_list { + } + backend_config_payload { + id: 0 + } + } + program_shape { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + id: 1 + root_id: 4294967300 + } + host_program_shape { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + entry_computation_id: 1 + schedule { + sequences { + key: 0 + value { + instruction_ids: 1 + instruction_ids: 0 + instruction_ids: 2 + } + } + sequences { + key: 1 + value { + instruction_ids: 4294967296 + instruction_ids: 4294967298 + instruction_ids: 4294967297 + instruction_ids: 4294967299 + instruction_ids: 4294967300 + } + } + } + input_output_alias { + } + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + buffer_donor { + } + frontend_attributes { + map { + key: "fingerprint_before_lhs" + value: "c36d8870be58d238935bebefbbc5b827" + } + } + payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" + } + config { + entry_computation_layout { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + replica_count: 2 + num_partitions: 1 + intra_op_parallelism_threads: -1 + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.01 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: false + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: false + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + } + static_device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + allow_spmd_sharding_propagation_to_parameters: false + sharding_config { + } + memory_fitting_level: EFFORT_O2 + schedule_config { + } + partition_size: 2 + } + } + buffer_assignment { + logical_buffers { + size: 4 + defined_at { + } + } + logical_buffers { + id: 1 + size: 4 + defined_at { + instruction_id: 1 + } + } + logical_buffers { + id: 2 + size: 4 + defined_at { + instruction_id: 2 + } + } + logical_buffers { + id: 10 + size: 32 + defined_at { + instruction_id: 4294967297 + } + } + logical_buffers { + id: 11 + size: 32 + defined_at { + instruction_id: 4294967298 + } + } + logical_buffers { + id: 12 + size: 4 + defined_at { + instruction_id: 4294967296 + } + } + logical_buffers { + id: 13 + size: 32 + defined_at { + instruction_id: 4294967299 + } + } + logical_buffers { + id: 14 + size: 16 + defined_at { + instruction_id: 4294967300 + } + } + buffer_allocations { + size: 32 + assigned { + logical_buffer_id: 10 + size: 32 + element_type: U32 + } + is_constant: true + } + buffer_allocations { + index: 1 + size: 32 + assigned { + logical_buffer_id: 11 + size: 32 + element_type: U32 + } + is_constant: true + } + buffer_allocations { + index: 2 + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 12 + size: 4 + element_type: U32 + } + assigned { + logical_buffer_id: 14 + size: 16 + element_type: U32 + } + } + buffer_allocations { + index: 3 + size: 4 + is_thread_local: true + assigned { + logical_buffer_id: 2 + size: 4 + element_type: U32 + } + } + buffer_allocations { + index: 4 + size: 4 + is_thread_local: true + assigned { + size: 4 + element_type: U32 + } + } + buffer_allocations { + index: 5 + size: 4 + is_thread_local: true + assigned { + logical_buffer_id: 1 + size: 4 + element_type: U32 + } + } + buffer_allocations { + index: 6 + size: 32 + assigned { + logical_buffer_id: 13 + size: 32 + element_type: U32 + } + } + heap_simulator_traces { + events { + buffer_id: 13 + computation_name: "main" + instruction_name: "loop_select_fusion" + } + events { + kind: FREE + buffer_id: 13 + computation_name: "main" + instruction_name: "loop_select_fusion" + } + whole_module_simulation: true + buffer_allocation_index: 3 + } + peak_buffers { + logical_buffer_ids: 10 + } + peak_buffers { + index: 1 + logical_buffer_ids: 11 + } + peak_buffers { + index: 2 + logical_buffer_ids: 14 + } + peak_buffers { + index: 3 + logical_buffer_ids: 2 + } + peak_buffers { + index: 4 + logical_buffer_ids: 0 + } + peak_buffers { + index: 5 + logical_buffer_ids: 1 + } + peak_buffers { + index: 6 + logical_buffer_ids: 13 + } + } + binary: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000p\013\000\000\000\000\000\000\360\006\000\000\000\000\000\000\nd\000\006@\0008\000\004\000@\000\022\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.nv.shared.reserved.0\000.nv.global\000.nv.constant4\000\000\000\000.nv.compat\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant4\000\000\000\000.nv.callgraph\000.nv.prototype\000.nv.merc.debug_frame\000.nv.merc.rela.nv.constant.pic\000.nv.merc.nv.constant.pic\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.global\000.nv.constant.pic\000.nv.compat\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant.pic\000.nv.callgraph\000.nv.prototype\000buffer_for_c1_1\000buffer_for_c0_1\000.nv.constant4\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\200\000\000\000 \240\013\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000<\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\034\001\000\000\021\000\014\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000,\001\000\000\021\000\014\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\010\000\000\000\000\000\000\000\002\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\002\000\001\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\200\000\000\000 \240\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\254\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\034\001\000\000\035 \014\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000,\001\000\000\035 \014\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\\\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\001\000\000\000\000\000\000J\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000(\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\216\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\350\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\350\003\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\234\004\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\007\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\203\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\004\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\306\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\260\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\004\000\000\000\000\000\0000\000\000\000\000\000\000\000\003\000\000\000\n\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000r\000\000\000\001\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \005\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0000\005\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000g\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\000\000\000@\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\342\000\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\350\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\367\000\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\000\000\0000\000\000\000\000\000\000\000\021\000\000\000\017\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\001\000\000}\000\000p\002\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000 \005\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000L\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\006\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000p\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000p\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000 \005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\0000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\000\000\000\000\000\000\000\020\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + gpu_compute_capability { + cuda_compute_capability { + major: 10 + feature_extension: ACCELERATED_FEATURES + } + } + module_name: "test" + program_shape { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + buffer_allocations { + values { + size: 32 + is_constant: true + } + values { + index: 1 + size: 32 + is_constant: true + } + values { + index: 2 + size: 16 + maybe_live_out: true + } + values { + index: 3 + size: 4 + is_thread_local: true + } + values { + index: 4 + size: 4 + is_thread_local: true + } + values { + index: 5 + size: 4 + is_thread_local: true + } + values { + index: 6 + size: 32 + } + } + output_info_map { + shape_index { + } + output_info { + allocation_index: 2 + } + } + constants { + symbol_name: "buffer_for_c1_1" + allocation_index: 1 + } + constants { + symbol_name: "buffer_for_c0_1" + } + thunks { + thunk_info { + profile_annotation: "id.1" + thunk_id: 1 + } + replica_id_thunk { + dest_buffer { + size: 4 + buffer_allocation_index: 2 + element_type: U32 + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_select_fusion" + thunk_id: 2 + } + custom_kernel_thunk { + args { + slice { + size: 32 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + args { + slice { + size: 32 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + args { + slice { + size: 4 + buffer_allocation_index: 2 + element_type: U32 + } + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + } + args { + slice { + size: 32 + buffer_allocation_index: 6 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + written: false + written: false + written: false + written: true + custom_kernel { + name: "loop_select_fusion" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\030\025\000\000\000\000\000\000\230\017\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_select_fusion\000.nv.capmerc.text.loop_select_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_select_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_select_fusion\000.nv.constant0.loop_select_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\244\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\323\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000/\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000K\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000^\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\004<\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\014\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\360\000\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000\003\031 \000\004\n\010\000\007\000\000\000\200\003 \000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\344\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000.\016\000\202{\004\377\000\340\000\000\000\n\000\000\000b\016\000\031y\t\000\000\000\000\000\000!\000\000\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\314\017\000\202{\006\377\000\346\000\000\000\n\000\000\000\342\016\000\201y\002\002\004\000\000\000\000\231\036\014\000$\037\000\014r\000\002\377\000\000\000pP\360\003\000\332\017\001\202\013\004\377\000\342\000\000\000\010\000\000\000\260.\000\202\013\005\377\000\343\000\000\000\010\000\000\000\244\016\000%x\004\t\004\000\000\000\004\000\216\007\000\314O\000\201y\005\004\004\000\000\000\000\231\036\014\000\242\016\000%x\006\t\004\000\000\000\006\000\216\007\000\312\217\000\206y\000\006\005\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\021\000\000@\200\271\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\004\n\370\000\004\000\000\000A\002\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\203\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\030\000\000\000\002\"\016\006\000\000B\000\000\000\001\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\002\"\016\006\000\000B\000\000\000A\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\0028\0162\370\000@\021\000\000\000\000\202\001\n\000\000\002\001@\001\000\000\000\000\000\000\000\000\000\000\000P\010\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\001\000\000\000\000\000\000\004\020\000\000\000\004\220\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\014\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\220\001\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\244\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\323\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\000\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000/\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000K\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\240\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\350\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000L\002\000\000\000\000\000\000\177\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\305\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\370\004\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\254\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000j\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\364\005\000\000\000\000\000\000\204\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\272\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\365\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\343\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\000\002\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\244\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\021\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000\240\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\240\014\000\000\000\000\000\000\016\001\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000V\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\256\r\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000k\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000 \016\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000|\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000D\016\000\000\000\000\000\000\224\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\330\016\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\272\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\360\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\330\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\360\016\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\030\025\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\030\025\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\000\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\003\000\000\000\000\000\000\240\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 4 + kernel_name: "loop_select_fusion" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 2 + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 3 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 8 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "ars.3" + thunk_id: 3 + } + reduce_scatter_thunk { + collective_config { + operand_element_type: U32 + } + buffers { + element_count: 8 + source_buffer { + slice { + size: 32 + buffer_allocation_index: 6 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + destination_buffer { + slice { + size: 16 + buffer_allocation_index: 2 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + reduction_kind: REDUCTION_KIND_SUM + } + } + executable_abi_version { + platform_name: "CUDA" + cuda_platform_version { + cuda_toolkit_version: "12.8.0" + cudnn_version: "9.10.2" + cub_version: "3.1.2" + } + } + buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 64B\n maybe_live_out allocation: 16B\n preallocated temp allocation: 32B\n preallocated temp fragmentation: 0B (0.00%)\n total allocation: 124B\n total fragmentation: 76B (61.29%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 32B\n\t\tXLA Label: fusion\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: reduce-scatter\n\t\tShape: u32[4]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 7:\n\t\tSize: 4B\n\t\tXLA Label: add\n\t\tShape: u32[]\n\t\t==========================\n" +} diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v2/AllGatherMixedTypes.pbtxt b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v2/AllGatherMixedTypes.pbtxt new file mode 100644 index 00000000000000..49c9f9b9b06505 --- /dev/null +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v2/AllGatherMixedTypes.pbtxt @@ -0,0 +1,2190 @@ +# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto +# proto-message: xla.HumanReadableAotExecutable +executable_and_options { + compile_options { + executable_build_options { + device_ordinal: -1 + result_layout { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.1 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: true + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: true + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + xla_gpu_enable_persistent_symmetric_memory: false + xla_gpu_experimental_enable_raft_for_stable_topk: false + xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND + xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET + } + num_replicas: 2 + num_partitions: 1 + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + comp_envs { + } + allow_spmd_sharding_propagation_to_parameters: false + process_count: 1 + memory_fitting_level: EFFORT_O2 + } + } + pjrt_client_name: "PjRtStreamExecutorClient" +} +gpu_executable { + hlo_module_with_config { + hlo_module { + name: "test" + entry_computation_name: "test_computation" + computations { + name: "fused_transpose" + instructions { + name: "param_0.1" + opcode: "parameter" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_0.1" + } + id: 4294967298 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.6.1" + opcode: "bitcast" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.6.1" + } + id: 4294967297 + operand_ids: 4294967298 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "transpose.1.1" + opcode: "transpose" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "transpose.1.1" + } + dimensions: 1 + dimensions: 0 + id: 4294967296 + operand_ids: 4294967297 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + result { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.1" + } + id: 1 + root_id: 4294967296 + is_fusion_computation: true + } + computations { + name: "fused_transpose.1" + instructions { + name: "param_0.3" + opcode: "parameter" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_0.3" + } + id: 8589934594 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.4.1" + opcode: "bitcast" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.4.1" + } + id: 8589934593 + operand_ids: 8589934594 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "transpose.3" + opcode: "transpose" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "transpose.3" + } + dimensions: 1 + dimensions: 0 + id: 8589934592 + operand_ids: 8589934593 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + result { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.3" + } + id: 2 + root_id: 8589934592 + is_fusion_computation: true + } + computations { + name: "fused_broadcast" + instructions { + name: "param_0.5" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "param_0.5" + } + id: 12884901890 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "p1.1.1" + opcode: "convert" + shape { + element_type: F32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "p1.1.1" + } + id: 12884901889 + operand_ids: 12884901890 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "broadcast.13.1" + opcode: "broadcast" + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "broadcast.13.1" + } + id: 12884901888 + operand_ids: 12884901889 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.5" + } + id: 3 + root_id: 12884901888 + is_fusion_computation: true + } + computations { + name: "wrapped_broadcast_computation" + instructions { + name: "param_0.7" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "param_0.7" + } + id: 21474836481 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "broadcast.12.1" + opcode: "broadcast" + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "broadcast.12.1" + } + id: 21474836480 + operand_ids: 21474836481 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.7" + } + id: 5 + root_id: 21474836480 + is_fusion_computation: true + } + computations { + name: "test_computation" + instructions { + name: "id.1" + opcode: "replica-id" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "id.1" + } + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "wrapped_broadcast" + opcode: "fusion" + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "wrapped_broadcast" + } + fusion_kind: "kLoop" + id: 9 + operand_ids: 0 + called_computation_ids: 5 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_broadcast_fusion" + opcode: "fusion" + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_broadcast_fusion" + } + fusion_kind: "kLoop" + id: 8 + operand_ids: 0 + called_computation_ids: 3 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "allgather.3" + opcode: "all-gather" + shape { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + metadata { + scheduling_name: "all-gather-start" + } + dimensions: 1 + id: 10 + operand_ids: 9 + operand_ids: 8 + frontend_attributes { + } + statistics_viz { + } + collective_device_list { + } + backend_config_payload { + id: 0 + } + } + instructions { + name: "ag0.1" + opcode: "get-tuple-element" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "ag0.1" + } + id: 1 + operand_ids: 10 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_transpose_fusion.1" + opcode: "fusion" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_transpose_fusion.1" + } + fusion_kind: "kLoop" + id: 7 + operand_ids: 1 + called_computation_ids: 2 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.24" + opcode: "bitcast" + shape { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.24" + } + id: 4 + operand_ids: 7 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "ag1.1" + opcode: "get-tuple-element" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "ag1.1" + } + tuple_index: 1 + id: 2 + operand_ids: 10 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_transpose_fusion" + opcode: "fusion" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_transpose_fusion" + } + fusion_kind: "kLoop" + id: 6 + operand_ids: 2 + called_computation_ids: 1 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.1.0" + opcode: "bitcast" + shape { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.1.0" + } + id: 5 + operand_ids: 6 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "out.1" + opcode: "tuple" + shape { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + metadata { + scheduling_name: "out.1" + } + id: 3 + operand_ids: 4 + operand_ids: 5 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + root_id: 3 + } + host_program_shape { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + schedule { + sequences { + key: 0 + value { + instruction_ids: 0 + instruction_ids: 8 + instruction_ids: 9 + instruction_ids: 10 + instruction_ids: 1 + instruction_ids: 7 + instruction_ids: 2 + instruction_ids: 6 + instruction_ids: 4 + instruction_ids: 5 + instruction_ids: 3 + } + } + } + input_output_alias { + } + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + buffer_donor { + } + frontend_attributes { + map { + key: "fingerprint_before_lhs" + value: "7b2e4f79c5524b5cbc566da178aa7fda" + } + } + payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" + } + config { + entry_computation_layout { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + replica_count: 2 + num_partitions: 1 + intra_op_parallelism_threads: -1 + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.1 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: true + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: true + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + xla_gpu_enable_persistent_symmetric_memory: false + xla_gpu_experimental_enable_raft_for_stable_topk: false + xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND + xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET + } + static_device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + allow_spmd_sharding_propagation_to_parameters: false + sharding_config { + } + memory_fitting_level: EFFORT_O2 + schedule_config { + } + partition_size: 2 + } + } + buffer_assignment { + logical_buffers { + id: 9 + size: 4 + defined_at { + } + } + logical_buffers { + id: 10 + size: 8 + defined_at { + instruction_id: 9 + } + } + logical_buffers { + id: 11 + size: 8 + defined_at { + instruction_id: 8 + } + } + logical_buffers { + id: 12 + size: 16 + defined_at { + instruction_id: 10 + } + } + logical_buffers { + id: 13 + size: 16 + defined_at { + shape_index: 0 + instruction_id: 10 + } + } + logical_buffers { + id: 14 + size: 16 + defined_at { + shape_index: 1 + instruction_id: 10 + } + } + logical_buffers { + id: 15 + size: 16 + defined_at { + instruction_id: 7 + } + } + logical_buffers { + id: 16 + size: 16 + defined_at { + instruction_id: 6 + } + } + logical_buffers { + id: 17 + size: 16 + defined_at { + instruction_id: 3 + } + } + buffer_allocations { + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 10 + size: 8 + element_type: U32 + } + assigned { + logical_buffer_id: 15 + size: 16 + element_type: U32 + } + } + buffer_allocations { + index: 1 + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 9 + size: 4 + element_type: U32 + } + assigned { + logical_buffer_id: 13 + size: 16 + element_type: U32 + } + assigned { + logical_buffer_id: 16 + size: 16 + element_type: F32 + } + } + buffer_allocations { + index: 2 + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 17 + size: 16 + element_type: TUPLE + } + is_tuple: true + } + buffer_allocations { + index: 3 + size: 520 + assigned { + logical_buffer_id: 11 + offset: 512 + size: 8 + element_type: F32 + } + assigned { + logical_buffer_id: 12 + size: 16 + element_type: TUPLE + } + assigned { + logical_buffer_id: 14 + offset: 256 + size: 16 + element_type: F32 + } + is_tuple: true + } + heap_simulator_traces { + events { + buffer_id: 11 + computation_name: "test_computation" + instruction_name: "loop_broadcast_fusion" + } + events { + buffer_id: 14 + computation_name: "test_computation" + instruction_name: "allgather.3" + } + events { + kind: FREE + buffer_id: 11 + computation_name: "test_computation" + instruction_name: "loop_broadcast_fusion" + } + events { + kind: FREE + buffer_id: 14 + computation_name: "test_computation" + instruction_name: "allgather.3" + } + whole_module_simulation: true + buffer_allocation_index: 3 + } + peak_buffers { + logical_buffer_ids: 15 + } + peak_buffers { + index: 1 + logical_buffer_ids: 16 + } + peak_buffers { + index: 2 + logical_buffer_ids: 17 + } + peak_buffers { + index: 3 + logical_buffer_ids: 12 + logical_buffer_ids: 11 + logical_buffer_ids: 14 + } + } + gpu_compute_capability { + cuda_compute_capability { + major: 10 + feature_extension: ACCELERATED_FEATURES + } + } + module_name: "test" + program_shape { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + buffer_allocations { + values { + size: 16 + maybe_live_out: true + } + values { + index: 1 + size: 16 + maybe_live_out: true + } + values { + index: 2 + size: 16 + maybe_live_out: true + is_tuple: true + } + values { + index: 3 + size: 520 + is_tuple: true + } + } + output_info_map { + shape_index { + } + output_info { + allocation_index: 2 + } + } + output_info_map { + shape_index { + indexes: 0 + } + output_info { + } + } + output_info_map { + shape_index { + indexes: 1 + } + output_info { + allocation_index: 1 + } + } + thunks { + thunk_info { + profile_annotation: "id.1" + thunk_id: 1 + } + replica_id_thunk { + dest_buffer { + size: 4 + buffer_allocation_index: 1 + element_type: U32 + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_broadcast_fusion" + thunk_id: 2 + } + custom_kernel_thunk { + args { + slice { + size: 4 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + } + args { + slice { + offset: 512 + size: 8 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "loop_broadcast_fusion" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_broadcast_fusion\000.nv.capmerc.text.loop_broadcast_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_broadcast_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_broadcast_fusion\000.nv.constant0.loop_broadcast_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000j\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004(\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\240\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\340\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000\"\016\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\342\000\000\000\n\000\000\000b\016\000\201y\002\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\342/\000Er\007\000\002\000\000\000\000\020 \000\000\312O\000\206y\000\004\007\000\000\000\004\031\020\014\000\342\017\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\014\000\000@\300\005\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\001\000\000\000\000\000\000\000\000\000\000\000\320\006\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\004\020\000\000\000\004\360\000\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\360\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\372\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000^\002\000\000\000\000\000\000\216\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\316\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\030\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\314\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000m\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\006\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\303\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\376\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\354\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\032\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000e\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000z\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\213\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\262\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\314\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\352\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "loop_broadcast_fusion" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 2 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "wrapped_broadcast" + thunk_id: 3 + } + custom_kernel_thunk { + args { + slice { + size: 4 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + } + args { + slice { + size: 8 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "wrapped_broadcast" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.wrapped_broadcast\000.nv.capmerc.text.wrapped_broadcast\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.wrapped_broadcast\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000wrapped_broadcast\000.nv.constant0.wrapped_broadcast\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\320\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\374\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000Z\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004$\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\220\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\340\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000\"\016\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\342\000\000\000\n\000\000\000b\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\312/\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\013\000\000@\300\002\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\007\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\000\000\000\000\000\000\000\004\020\000\000\000\004\340\000\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\340\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\320\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\374\000\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\360\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\342\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000F\002\000\000\000\000\000\000z\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\302\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\350\004\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\234\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000i\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\344\005\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\267\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\362\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000X\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\340\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000f\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000w\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\232\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\264\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\322\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "wrapped_broadcast" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 2 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "allgather.3" + thunk_id: 4 + } + all_gather_thunk { + collective_config { + operand_element_type: U32 + operand_element_type: F32 + } + buffers { + element_count: 2 + source_buffer { + slice { + size: 8 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + destination_buffer { + slice { + size: 16 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + } + buffers { + element_count: 2 + source_buffer { + slice { + offset: 512 + size: 8 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + destination_buffer { + slice { + offset: 256 + size: 16 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_transpose_fusion.1" + thunk_id: 5 + } + custom_kernel_thunk { + args { + slice { + size: 16 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + args { + slice { + size: 16 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "loop_transpose_fusion_1" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion_1\000.nv.capmerc.text.loop_transpose_fusion_1\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_transpose_fusion_1\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion_1\000.nv.constant0.loop_transpose_fusion_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\342\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Z\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000r\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\340\000\000\000\n\000\000\000b\016\000\254w\004\377\000k\000\000\000\n\000\010\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\342\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242N\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\017\000\000@\340-\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\005\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\004\020\000\000\000\004`\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000`\001\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\342\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Z\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\006\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000j\002\000\000\000\000\000\000\230\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\004\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\324\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\0000\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\344\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\006\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000o\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\006\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\311\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\001\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\362\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000F\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000o\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\204\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\225\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\276\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\330\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "loop_transpose_fusion_1" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 4 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_transpose_fusion" + thunk_id: 6 + } + custom_kernel_thunk { + args { + slice { + offset: 256 + size: 16 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + args { + slice { + size: 16 + buffer_allocation_index: 1 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "loop_transpose_fusion" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion\000.nv.capmerc.text.loop_transpose_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_transpose_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion\000.nv.constant0.loop_transpose_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000j\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\340\000\000\000\n\000\000\000b\016\000\254w\004\377\000k\000\000\000\n\000\010\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\342\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242N\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\017\000\000@\340-\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\005\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\004\020\000\000\000\004`\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000`\001\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\372\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000^\002\000\000\000\000\000\000\216\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\316\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\030\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\314\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000m\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\006\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\303\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\376\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\354\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\032\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000e\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000z\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\213\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\262\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\314\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\352\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "loop_transpose_fusion" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 4 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + executable_abi_version { + platform_name: "CUDA" + cuda_platform_version { + cuda_toolkit_version: "12.8.0" + cudnn_version: "9.10.2" + cub_version: "3.1.2" + } + } + buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 0B\n maybe_live_out allocation: 48B\n preallocated temp allocation: 520B\n preallocated temp fragmentation: 240B (46.15%)\n total allocation: 568B\n total fragmentation: 504B (88.73%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: (u32[2,2], f32[2,2])\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 16B\n\t\tXLA Label: tuple\n\t\tShape: (u32[4], f32[4])\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: u32[2,2]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 8B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,1]\n\t\t==========================\n" +} diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v2/ReduceScatter.pbtxt b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v2/ReduceScatter.pbtxt new file mode 100644 index 00000000000000..07c084a61d3c00 --- /dev/null +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_b200/v2/ReduceScatter.pbtxt @@ -0,0 +1,1624 @@ +# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto +# proto-message: xla.HumanReadableAotExecutable +executable_and_options { + compile_options { + executable_build_options { + device_ordinal: -1 + result_layout { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.1 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: true + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: true + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + xla_gpu_enable_persistent_symmetric_memory: false + xla_gpu_experimental_enable_raft_for_stable_topk: false + xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND + xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET + } + num_replicas: 2 + num_partitions: 1 + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + comp_envs { + } + allow_spmd_sharding_propagation_to_parameters: false + process_count: 1 + memory_fitting_level: EFFORT_O2 + } + } + pjrt_client_name: "PjRtStreamExecutorClient" +} +gpu_executable { + hlo_module_with_config { + hlo_module { + name: "test" + entry_computation_name: "main" + computations { + name: "add" + instructions { + name: "lhs" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "lhs" + } + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "rhs" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "rhs" + } + parameter_number: 1 + id: 1 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "add.1" + opcode: "add" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "add.1" + } + id: 2 + operand_ids: 0 + operand_ids: 1 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + parameter_names: "lhs" + parameter_names: "rhs" + } + root_id: 2 + } + computations { + name: "fused_select" + instructions { + name: "param_2.2" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "param_2.2" + } + parameter_number: 2 + id: 8589934597 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "zero_2" + opcode: "constant" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + } + literal { + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + u32s: 0 + } + id: 8589934598 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "p.2" + opcode: "compare" + shape { + element_type: PRED + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "p.2" + } + id: 8589934596 + operand_ids: 8589934597 + operand_ids: 8589934598 + comparison_direction: "EQ" + frontend_attributes { + } + statistics_viz { + } + comparison_order: "TOTAL" + } + instructions { + name: "pb.2" + opcode: "broadcast" + shape { + element_type: PRED + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "pb.2" + } + id: 8589934593 + operand_ids: 8589934596 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "param_0.1" + opcode: "parameter" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_0.1" + } + id: 8589934594 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "param_1.1" + opcode: "parameter" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_1.1" + } + parameter_number: 1 + id: 8589934595 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "data.2" + opcode: "select" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "data.2" + } + id: 8589934592 + operand_ids: 8589934593 + operand_ids: 8589934594 + operand_ids: 8589934595 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + parameters { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + parameter_names: "param_0.1" + parameter_names: "param_1.1" + parameter_names: "param_2.2" + } + id: 2 + root_id: 8589934592 + is_fusion_computation: true + } + computations { + name: "main" + instructions { + name: "c0_1" + opcode: "constant" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + } + literal { + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + u32s: 1 + u32s: 2 + u32s: 3 + u32s: 4 + u32s: 5 + u32s: 6 + u32s: 7 + u32s: 8 + } + id: 4294967297 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "c1_1" + opcode: "constant" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + } + literal { + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + u32s: 10 + u32s: 11 + u32s: 12 + u32s: 13 + u32s: 14 + u32s: 15 + u32s: 16 + u32s: 17 + } + id: 4294967298 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "id.1" + opcode: "replica-id" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "id.1" + } + id: 4294967296 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_select_fusion" + opcode: "fusion" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_select_fusion" + } + fusion_kind: "kLoop" + id: 4294967299 + operand_ids: 4294967297 + operand_ids: 4294967298 + operand_ids: 4294967296 + called_computation_ids: 2 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "ars.3" + opcode: "reduce-scatter" + shape { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "reduce-scatter-start" + } + dimensions: 0 + id: 4294967300 + operand_ids: 4294967299 + called_computation_ids: 0 + frontend_attributes { + } + statistics_viz { + } + collective_device_list { + } + backend_config_payload { + id: 0 + } + } + program_shape { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + id: 1 + root_id: 4294967300 + } + host_program_shape { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + entry_computation_id: 1 + schedule { + sequences { + key: 0 + value { + instruction_ids: 1 + instruction_ids: 0 + instruction_ids: 2 + } + } + sequences { + key: 1 + value { + instruction_ids: 4294967296 + instruction_ids: 4294967298 + instruction_ids: 4294967297 + instruction_ids: 4294967299 + instruction_ids: 4294967300 + } + } + } + input_output_alias { + } + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + buffer_donor { + } + frontend_attributes { + map { + key: "fingerprint_before_lhs" + value: "c36d8870be58d238935bebefbbc5b827" + } + } + payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" + } + config { + entry_computation_layout { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + replica_count: 2 + num_partitions: 1 + intra_op_parallelism_threads: -1 + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.1 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: true + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: true + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + xla_gpu_enable_persistent_symmetric_memory: false + xla_gpu_experimental_enable_raft_for_stable_topk: false + xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND + xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET + } + static_device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + allow_spmd_sharding_propagation_to_parameters: false + sharding_config { + } + memory_fitting_level: EFFORT_O2 + schedule_config { + } + partition_size: 2 + } + } + buffer_assignment { + logical_buffers { + size: 4 + defined_at { + } + } + logical_buffers { + id: 1 + size: 4 + defined_at { + instruction_id: 1 + } + } + logical_buffers { + id: 2 + size: 4 + defined_at { + instruction_id: 2 + } + } + logical_buffers { + id: 10 + size: 32 + defined_at { + instruction_id: 4294967297 + } + } + logical_buffers { + id: 11 + size: 32 + defined_at { + instruction_id: 4294967298 + } + } + logical_buffers { + id: 12 + size: 4 + defined_at { + instruction_id: 4294967296 + } + } + logical_buffers { + id: 13 + size: 32 + defined_at { + instruction_id: 4294967299 + } + } + logical_buffers { + id: 14 + size: 16 + defined_at { + instruction_id: 4294967300 + } + } + buffer_allocations { + size: 32 + assigned { + logical_buffer_id: 10 + size: 32 + element_type: U32 + } + is_constant: true + } + buffer_allocations { + index: 1 + size: 32 + assigned { + logical_buffer_id: 11 + size: 32 + element_type: U32 + } + is_constant: true + } + buffer_allocations { + index: 2 + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 12 + size: 4 + element_type: U32 + } + assigned { + logical_buffer_id: 14 + size: 16 + element_type: U32 + } + } + buffer_allocations { + index: 3 + size: 4 + is_thread_local: true + assigned { + logical_buffer_id: 2 + size: 4 + element_type: U32 + } + } + buffer_allocations { + index: 4 + size: 4 + is_thread_local: true + assigned { + size: 4 + element_type: U32 + } + } + buffer_allocations { + index: 5 + size: 4 + is_thread_local: true + assigned { + logical_buffer_id: 1 + size: 4 + element_type: U32 + } + } + buffer_allocations { + index: 6 + size: 32 + assigned { + logical_buffer_id: 13 + size: 32 + element_type: U32 + } + } + heap_simulator_traces { + events { + buffer_id: 13 + computation_name: "main" + instruction_name: "loop_select_fusion" + } + events { + kind: FREE + buffer_id: 13 + computation_name: "main" + instruction_name: "loop_select_fusion" + } + whole_module_simulation: true + buffer_allocation_index: 3 + } + peak_buffers { + logical_buffer_ids: 10 + } + peak_buffers { + index: 1 + logical_buffer_ids: 11 + } + peak_buffers { + index: 2 + logical_buffer_ids: 14 + } + peak_buffers { + index: 3 + logical_buffer_ids: 2 + } + peak_buffers { + index: 4 + logical_buffer_ids: 0 + } + peak_buffers { + index: 5 + logical_buffer_ids: 1 + } + peak_buffers { + index: 6 + logical_buffer_ids: 13 + } + } + binary: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000p\013\000\000\000\000\000\000\360\006\000\000\000\000\000\000\nd\000\006@\0008\000\004\000@\000\022\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.nv.shared.reserved.0\000.nv.global\000.nv.constant4\000\000\000\000.nv.compat\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant4\000\000\000\000.nv.callgraph\000.nv.prototype\000.nv.merc.debug_frame\000.nv.merc.rela.nv.constant.pic\000.nv.merc.nv.constant.pic\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.global\000.nv.constant.pic\000.nv.compat\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant.pic\000.nv.callgraph\000.nv.prototype\000buffer_for_c1_1\000buffer_for_c0_1\000.nv.constant4\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\200\000\000\000 \240\013\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000<\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\034\001\000\000\021\000\014\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000,\001\000\000\021\000\014\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\010\000\000\000\000\000\000\000\002\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\002\000\001\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\200\000\000\000 \240\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\254\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\034\001\000\000\035 \014\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000,\001\000\000\035 \014\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\\\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\001\000\000\000\000\000\000J\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000(\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\216\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\350\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\350\003\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\234\004\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\007\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\203\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\004\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\306\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\260\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\004\000\000\000\000\000\0000\000\000\000\000\000\000\000\003\000\000\000\n\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000r\000\000\000\001\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \005\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0000\005\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000g\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\000\000\000@\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\342\000\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\350\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\367\000\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\000\000\0000\000\000\000\000\000\000\000\021\000\000\000\017\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\001\000\000}\000\000p\002\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000 \005\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000L\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\006\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000p\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000p\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000 \005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\0000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\000\000\000\000\000\000\000\020\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + gpu_compute_capability { + cuda_compute_capability { + major: 10 + feature_extension: ACCELERATED_FEATURES + } + } + module_name: "test" + program_shape { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + buffer_allocations { + values { + size: 32 + is_constant: true + } + values { + index: 1 + size: 32 + is_constant: true + } + values { + index: 2 + size: 16 + maybe_live_out: true + } + values { + index: 3 + size: 4 + is_thread_local: true + } + values { + index: 4 + size: 4 + is_thread_local: true + } + values { + index: 5 + size: 4 + is_thread_local: true + } + values { + index: 6 + size: 32 + } + } + output_info_map { + shape_index { + } + output_info { + allocation_index: 2 + } + } + constants { + symbol_name: "buffer_for_c1_1" + allocation_index: 1 + } + constants { + symbol_name: "buffer_for_c0_1" + } + thunks { + thunk_info { + profile_annotation: "id.1" + thunk_id: 1 + } + replica_id_thunk { + dest_buffer { + size: 4 + buffer_allocation_index: 2 + element_type: U32 + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_select_fusion" + thunk_id: 2 + } + custom_kernel_thunk { + args { + slice { + size: 32 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + args { + slice { + size: 32 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + args { + slice { + size: 4 + buffer_allocation_index: 2 + element_type: U32 + } + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + } + args { + slice { + size: 32 + buffer_allocation_index: 6 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + written: false + written: false + written: false + written: true + custom_kernel { + name: "loop_select_fusion" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\030\025\000\000\000\000\000\000\230\017\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_select_fusion\000.nv.capmerc.text.loop_select_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_select_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_select_fusion\000.nv.constant0.loop_select_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\244\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\323\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000/\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000K\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000^\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\004<\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\014\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\360\000\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000\003\031 \000\004\n\010\000\007\000\000\000\200\003 \000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\344\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000.\016\000\202{\004\377\000\340\000\000\000\n\000\000\000b\016\000\031y\t\000\000\000\000\000\000!\000\000\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\314\017\000\202{\006\377\000\346\000\000\000\n\000\000\000\342\016\000\201y\002\002\004\000\000\000\000\231\036\014\000$\037\000\014r\000\002\377\000\000\000pP\360\003\000\332\017\001\202\013\004\377\000\342\000\000\000\010\000\000\000\260.\000\202\013\005\377\000\343\000\000\000\010\000\000\000\244\016\000%x\004\t\004\000\000\000\004\000\216\007\000\314O\000\201y\005\004\004\000\000\000\000\231\036\014\000\242\016\000%x\006\t\004\000\000\000\006\000\216\007\000\312\217\000\206y\000\006\005\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\021\000\000@\200\271\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\004\n\370\000\004\000\000\000A\002\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\203\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\030\000\000\000\002\"\016\006\000\000B\000\000\000\001\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\002\"\016\006\000\000B\000\000\000A\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\0028\0162\370\000@\021\000\000\000\000\202\001\n\000\000\002\001@\001\000\000\000\000\000\000\000\000\000\000\000P\010\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\001\000\000\000\000\000\000\004\020\000\000\000\004\220\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\014\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\220\001\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\244\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\323\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\000\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000/\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000K\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\240\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\350\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000L\002\000\000\000\000\000\000\177\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\305\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\370\004\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\254\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000j\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\364\005\000\000\000\000\000\000\204\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\272\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\365\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\343\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\000\002\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\244\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\021\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000\240\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\240\014\000\000\000\000\000\000\016\001\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000V\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\256\r\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000k\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000 \016\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000|\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000D\016\000\000\000\000\000\000\224\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\330\016\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\272\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\360\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\330\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\360\016\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\030\025\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\030\025\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\000\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\003\000\000\000\000\000\000\240\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 4 + kernel_name: "loop_select_fusion" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 2 + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 3 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 8 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "ars.3" + thunk_id: 3 + } + reduce_scatter_thunk { + collective_config { + operand_element_type: U32 + } + buffers { + element_count: 8 + source_buffer { + slice { + size: 32 + buffer_allocation_index: 6 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + destination_buffer { + slice { + size: 16 + buffer_allocation_index: 2 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + reduction_kind: REDUCTION_KIND_SUM + } + } + executable_abi_version { + platform_name: "CUDA" + cuda_platform_version { + cuda_toolkit_version: "12.8.0" + cudnn_version: "9.10.2" + cub_version: "3.1.2" + } + } + buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 64B\n maybe_live_out allocation: 16B\n preallocated temp allocation: 32B\n preallocated temp fragmentation: 0B (0.00%)\n total allocation: 124B\n total fragmentation: 76B (61.29%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 32B\n\t\tXLA Label: fusion\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: reduce-scatter\n\t\tShape: u32[4]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 7:\n\t\tSize: 4B\n\t\tXLA Label: add\n\t\tShape: u32[]\n\t\t==========================\n" +} diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v1/AllGatherMixedTypes.pbtxt b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v1/AllGatherMixedTypes.pbtxt new file mode 100644 index 00000000000000..78bf6596da5d5a --- /dev/null +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v1/AllGatherMixedTypes.pbtxt @@ -0,0 +1,2180 @@ +# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto +# proto-message: xla.HumanReadableAotExecutable +executable_and_options { + compile_options { + executable_build_options { + device_ordinal: -1 + result_layout { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.01 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "NEON" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: false + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: false + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + } + num_replicas: 2 + num_partitions: 1 + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + comp_envs { + } + allow_spmd_sharding_propagation_to_parameters: false + process_count: 1 + memory_fitting_level: EFFORT_O2 + } + } + pjrt_client_name: "PjRtStreamExecutorClient" +} +gpu_executable { + hlo_module_with_config { + hlo_module { + name: "test" + entry_computation_name: "test_computation" + computations { + name: "fused_transpose" + instructions { + name: "param_0.1" + opcode: "parameter" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_0.1" + } + id: 4294967298 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.6.1" + opcode: "bitcast" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.6.1" + } + id: 4294967297 + operand_ids: 4294967298 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "transpose.1.1" + opcode: "transpose" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "transpose.1.1" + } + dimensions: 1 + dimensions: 0 + id: 4294967296 + operand_ids: 4294967297 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + result { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.1" + } + id: 1 + root_id: 4294967296 + is_fusion_computation: true + } + computations { + name: "fused_transpose.1" + instructions { + name: "param_0.3" + opcode: "parameter" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_0.3" + } + id: 8589934594 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.4.1" + opcode: "bitcast" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.4.1" + } + id: 8589934593 + operand_ids: 8589934594 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "transpose.3" + opcode: "transpose" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "transpose.3" + } + dimensions: 1 + dimensions: 0 + id: 8589934592 + operand_ids: 8589934593 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + result { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.3" + } + id: 2 + root_id: 8589934592 + is_fusion_computation: true + } + computations { + name: "fused_broadcast" + instructions { + name: "param_0.5" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "param_0.5" + } + id: 12884901890 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "p1.1.1" + opcode: "convert" + shape { + element_type: F32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "p1.1.1" + } + id: 12884901889 + operand_ids: 12884901890 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "broadcast.13.1" + opcode: "broadcast" + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "broadcast.13.1" + } + id: 12884901888 + operand_ids: 12884901889 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.5" + } + id: 3 + root_id: 12884901888 + is_fusion_computation: true + } + computations { + name: "wrapped_broadcast_computation" + instructions { + name: "param_0.7" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "param_0.7" + } + id: 21474836481 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "broadcast.12.1" + opcode: "broadcast" + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "broadcast.12.1" + } + id: 21474836480 + operand_ids: 21474836481 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.7" + } + id: 5 + root_id: 21474836480 + is_fusion_computation: true + } + computations { + name: "test_computation" + instructions { + name: "id.1" + opcode: "replica-id" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "id.1" + } + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "wrapped_broadcast" + opcode: "fusion" + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "wrapped_broadcast" + } + fusion_kind: "kLoop" + id: 9 + operand_ids: 0 + called_computation_ids: 5 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_broadcast_fusion" + opcode: "fusion" + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_broadcast_fusion" + } + fusion_kind: "kLoop" + id: 8 + operand_ids: 0 + called_computation_ids: 3 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "allgather.3" + opcode: "all-gather" + shape { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + metadata { + scheduling_name: "all-gather-start" + } + dimensions: 1 + id: 10 + operand_ids: 9 + operand_ids: 8 + frontend_attributes { + } + statistics_viz { + } + collective_device_list { + } + backend_config_payload { + id: 0 + } + } + instructions { + name: "ag0.1" + opcode: "get-tuple-element" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "ag0.1" + } + id: 1 + operand_ids: 10 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_transpose_fusion.1" + opcode: "fusion" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_transpose_fusion.1" + } + fusion_kind: "kLoop" + id: 7 + operand_ids: 1 + called_computation_ids: 2 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.24" + opcode: "bitcast" + shape { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.24" + } + id: 4 + operand_ids: 7 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "ag1.1" + opcode: "get-tuple-element" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "ag1.1" + } + tuple_index: 1 + id: 2 + operand_ids: 10 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_transpose_fusion" + opcode: "fusion" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_transpose_fusion" + } + fusion_kind: "kLoop" + id: 6 + operand_ids: 2 + called_computation_ids: 1 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.1.0" + opcode: "bitcast" + shape { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.1.0" + } + id: 5 + operand_ids: 6 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "out.1" + opcode: "tuple" + shape { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + metadata { + scheduling_name: "out.1" + } + id: 3 + operand_ids: 4 + operand_ids: 5 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + root_id: 3 + } + host_program_shape { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + schedule { + sequences { + key: 0 + value { + instruction_ids: 0 + instruction_ids: 8 + instruction_ids: 9 + instruction_ids: 10 + instruction_ids: 1 + instruction_ids: 7 + instruction_ids: 2 + instruction_ids: 6 + instruction_ids: 4 + instruction_ids: 5 + instruction_ids: 3 + } + } + } + input_output_alias { + } + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + buffer_donor { + } + frontend_attributes { + map { + key: "fingerprint_before_lhs" + value: "7b2e4f79c5524b5cbc566da178aa7fda" + } + } + payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" + } + config { + entry_computation_layout { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + replica_count: 2 + num_partitions: 1 + intra_op_parallelism_threads: -1 + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.01 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "NEON" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: false + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: false + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + } + static_device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + allow_spmd_sharding_propagation_to_parameters: false + sharding_config { + } + memory_fitting_level: EFFORT_O2 + schedule_config { + } + partition_size: 2 + } + } + buffer_assignment { + logical_buffers { + id: 9 + size: 4 + defined_at { + } + } + logical_buffers { + id: 10 + size: 8 + defined_at { + instruction_id: 9 + } + } + logical_buffers { + id: 11 + size: 8 + defined_at { + instruction_id: 8 + } + } + logical_buffers { + id: 12 + size: 16 + defined_at { + instruction_id: 10 + } + } + logical_buffers { + id: 13 + size: 16 + defined_at { + shape_index: 0 + instruction_id: 10 + } + } + logical_buffers { + id: 14 + size: 16 + defined_at { + shape_index: 1 + instruction_id: 10 + } + } + logical_buffers { + id: 15 + size: 16 + defined_at { + instruction_id: 7 + } + } + logical_buffers { + id: 16 + size: 16 + defined_at { + instruction_id: 6 + } + } + logical_buffers { + id: 17 + size: 16 + defined_at { + instruction_id: 3 + } + } + buffer_allocations { + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 10 + size: 8 + element_type: U32 + } + assigned { + logical_buffer_id: 15 + size: 16 + element_type: U32 + } + } + buffer_allocations { + index: 1 + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 9 + size: 4 + element_type: U32 + } + assigned { + logical_buffer_id: 13 + size: 16 + element_type: U32 + } + assigned { + logical_buffer_id: 16 + size: 16 + element_type: F32 + } + } + buffer_allocations { + index: 2 + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 17 + size: 16 + element_type: TUPLE + } + is_tuple: true + } + buffer_allocations { + index: 3 + size: 520 + assigned { + logical_buffer_id: 11 + offset: 512 + size: 8 + element_type: F32 + } + assigned { + logical_buffer_id: 12 + size: 16 + element_type: TUPLE + } + assigned { + logical_buffer_id: 14 + offset: 256 + size: 16 + element_type: F32 + } + is_tuple: true + } + heap_simulator_traces { + events { + buffer_id: 11 + computation_name: "test_computation" + instruction_name: "loop_broadcast_fusion" + } + events { + buffer_id: 14 + computation_name: "test_computation" + instruction_name: "allgather.3" + } + events { + kind: FREE + buffer_id: 11 + computation_name: "test_computation" + instruction_name: "loop_broadcast_fusion" + } + events { + kind: FREE + buffer_id: 14 + computation_name: "test_computation" + instruction_name: "allgather.3" + } + whole_module_simulation: true + buffer_allocation_index: 3 + } + peak_buffers { + logical_buffer_ids: 15 + } + peak_buffers { + index: 1 + logical_buffer_ids: 16 + } + peak_buffers { + index: 2 + logical_buffer_ids: 17 + } + peak_buffers { + index: 3 + logical_buffer_ids: 12 + logical_buffer_ids: 11 + logical_buffer_ids: 14 + } + } + gpu_compute_capability { + cuda_compute_capability { + major: 10 + feature_extension: ACCELERATED_FEATURES + } + } + module_name: "test" + program_shape { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + buffer_allocations { + values { + size: 16 + maybe_live_out: true + } + values { + index: 1 + size: 16 + maybe_live_out: true + } + values { + index: 2 + size: 16 + maybe_live_out: true + is_tuple: true + } + values { + index: 3 + size: 520 + is_tuple: true + } + } + output_info_map { + shape_index { + } + output_info { + allocation_index: 2 + } + } + output_info_map { + shape_index { + indexes: 0 + } + output_info { + } + } + output_info_map { + shape_index { + indexes: 1 + } + output_info { + allocation_index: 1 + } + } + thunks { + thunk_info { + profile_annotation: "id.1" + thunk_id: 1 + } + replica_id_thunk { + dest_buffer { + size: 4 + buffer_allocation_index: 1 + element_type: U32 + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_broadcast_fusion" + thunk_id: 2 + } + custom_kernel_thunk { + args { + slice { + size: 4 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + } + args { + slice { + offset: 512 + size: 8 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "loop_broadcast_fusion" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_broadcast_fusion\000.nv.capmerc.text.loop_broadcast_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_broadcast_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_broadcast_fusion\000.nv.constant0.loop_broadcast_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000j\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004(\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\240\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\340\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000\"\016\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\342\000\000\000\n\000\000\000b\016\000\201y\002\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\342/\000Er\007\000\002\000\000\000\000\020 \000\000\312O\000\206y\000\004\007\000\000\000\004\031\020\014\000\342\017\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\014\000\000@\300\005\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\001\000\000\000\000\000\000\000\000\000\000\000\320\006\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\004\020\000\000\000\004\360\000\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\360\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\372\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000^\002\000\000\000\000\000\000\216\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\316\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\030\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\314\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000m\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\006\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\303\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\376\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\354\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\032\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000e\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000z\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\213\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\262\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\314\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\352\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "loop_broadcast_fusion" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 2 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "wrapped_broadcast" + thunk_id: 3 + } + custom_kernel_thunk { + args { + slice { + size: 4 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + } + args { + slice { + size: 8 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "wrapped_broadcast" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.wrapped_broadcast\000.nv.capmerc.text.wrapped_broadcast\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.wrapped_broadcast\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000wrapped_broadcast\000.nv.constant0.wrapped_broadcast\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\320\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\374\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000Z\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004$\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\220\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\340\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000\"\016\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\342\000\000\000\n\000\000\000b\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\312/\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\013\000\000@\300\002\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\007\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\000\000\000\000\000\000\000\004\020\000\000\000\004\340\000\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\340\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\320\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\374\000\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\360\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\342\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000F\002\000\000\000\000\000\000z\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\302\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\350\004\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\234\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000i\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\344\005\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\267\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\362\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000X\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\340\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000f\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000w\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\232\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\264\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\322\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "wrapped_broadcast" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 2 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "allgather.3" + thunk_id: 4 + } + all_gather_thunk { + collective_config { + operand_element_type: U32 + operand_element_type: F32 + } + buffers { + element_count: 2 + source_buffer { + slice { + size: 8 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + destination_buffer { + slice { + size: 16 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + } + buffers { + element_count: 2 + source_buffer { + slice { + offset: 512 + size: 8 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + destination_buffer { + slice { + offset: 256 + size: 16 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_transpose_fusion.1" + thunk_id: 5 + } + custom_kernel_thunk { + args { + slice { + size: 16 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + args { + slice { + size: 16 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "loop_transpose_fusion_1" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion_1\000.nv.capmerc.text.loop_transpose_fusion_1\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_transpose_fusion_1\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion_1\000.nv.constant0.loop_transpose_fusion_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\342\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Z\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000r\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\340\000\000\000\n\000\000\000b\016\000\254w\004\377\000k\000\000\000\n\000\010\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\342\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242N\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\017\000\000@\340-\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\005\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\004\020\000\000\000\004`\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000`\001\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\342\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Z\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\006\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000j\002\000\000\000\000\000\000\230\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\004\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\324\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\0000\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\344\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\006\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000o\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\006\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\311\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\001\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\362\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000F\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000o\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\204\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\225\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\276\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\330\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "loop_transpose_fusion_1" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 4 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_transpose_fusion" + thunk_id: 6 + } + custom_kernel_thunk { + args { + slice { + offset: 256 + size: 16 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + args { + slice { + size: 16 + buffer_allocation_index: 1 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "loop_transpose_fusion" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion\000.nv.capmerc.text.loop_transpose_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_transpose_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion\000.nv.constant0.loop_transpose_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000j\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\340\000\000\000\n\000\000\000b\016\000\254w\004\377\000k\000\000\000\n\000\010\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\342\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242N\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\017\000\000@\340-\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\005\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\004\020\000\000\000\004`\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000`\001\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\372\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000^\002\000\000\000\000\000\000\216\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\316\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\030\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\314\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000m\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\006\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\303\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\376\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\354\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\032\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000e\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000z\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\213\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\262\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\314\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\352\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "loop_transpose_fusion" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 4 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + executable_abi_version { + platform_name: "CUDA" + cuda_platform_version { + cuda_toolkit_version: "12.8.0" + cudnn_version: "9.10.2" + cub_version: "3.1.2" + } + } + buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 0B\n maybe_live_out allocation: 48B\n preallocated temp allocation: 520B\n preallocated temp fragmentation: 240B (46.15%)\n total allocation: 568B\n total fragmentation: 504B (88.73%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: (u32[2,2], f32[2,2])\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 16B\n\t\tXLA Label: tuple\n\t\tShape: (u32[4], f32[4])\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: u32[2,2]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 8B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,1]\n\t\t==========================\n" +} diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v1/ReduceScatter.pbtxt b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v1/ReduceScatter.pbtxt new file mode 100644 index 00000000000000..69429709332001 --- /dev/null +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v1/ReduceScatter.pbtxt @@ -0,0 +1,1615 @@ +# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto +# proto-message: xla.HumanReadableAotExecutable +executable_and_options { + compile_options { + executable_build_options { + device_ordinal: -1 + result_layout { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.01 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "NEON" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: false + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: false + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + } + num_replicas: 2 + num_partitions: 1 + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + comp_envs { + } + allow_spmd_sharding_propagation_to_parameters: false + process_count: 1 + memory_fitting_level: EFFORT_O2 + } + } + pjrt_client_name: "PjRtStreamExecutorClient" +} +gpu_executable { + hlo_module_with_config { + hlo_module { + name: "test" + entry_computation_name: "main" + computations { + name: "add" + instructions { + name: "lhs" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "lhs" + } + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "rhs" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "rhs" + } + parameter_number: 1 + id: 1 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "add.1" + opcode: "add" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "add.1" + } + id: 2 + operand_ids: 0 + operand_ids: 1 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + parameter_names: "lhs" + parameter_names: "rhs" + } + root_id: 2 + } + computations { + name: "fused_select" + instructions { + name: "param_2.2" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "param_2.2" + } + parameter_number: 2 + id: 8589934597 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "zero_2" + opcode: "constant" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + } + literal { + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + u32s: 0 + } + id: 8589934598 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "p.2" + opcode: "compare" + shape { + element_type: PRED + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "p.2" + } + id: 8589934596 + operand_ids: 8589934597 + operand_ids: 8589934598 + comparison_direction: "EQ" + frontend_attributes { + } + comparison_type: "UNSIGNED" + statistics_viz { + } + comparison_order: "TOTAL" + } + instructions { + name: "pb.2" + opcode: "broadcast" + shape { + element_type: PRED + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "pb.2" + } + id: 8589934593 + operand_ids: 8589934596 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "param_0.1" + opcode: "parameter" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_0.1" + } + id: 8589934594 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "param_1.1" + opcode: "parameter" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_1.1" + } + parameter_number: 1 + id: 8589934595 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "data.2" + opcode: "select" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "data.2" + } + id: 8589934592 + operand_ids: 8589934593 + operand_ids: 8589934594 + operand_ids: 8589934595 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + parameters { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + parameter_names: "param_0.1" + parameter_names: "param_1.1" + parameter_names: "param_2.2" + } + id: 2 + root_id: 8589934592 + is_fusion_computation: true + } + computations { + name: "main" + instructions { + name: "c0_1" + opcode: "constant" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + } + literal { + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + u32s: 1 + u32s: 2 + u32s: 3 + u32s: 4 + u32s: 5 + u32s: 6 + u32s: 7 + u32s: 8 + } + id: 4294967297 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "c1_1" + opcode: "constant" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + } + literal { + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + u32s: 10 + u32s: 11 + u32s: 12 + u32s: 13 + u32s: 14 + u32s: 15 + u32s: 16 + u32s: 17 + } + id: 4294967298 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "id.1" + opcode: "replica-id" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "id.1" + } + id: 4294967296 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_select_fusion" + opcode: "fusion" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_select_fusion" + } + fusion_kind: "kLoop" + id: 4294967299 + operand_ids: 4294967297 + operand_ids: 4294967298 + operand_ids: 4294967296 + called_computation_ids: 2 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "ars.3" + opcode: "reduce-scatter" + shape { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "reduce-scatter-start" + } + dimensions: 0 + id: 4294967300 + operand_ids: 4294967299 + called_computation_ids: 0 + frontend_attributes { + } + statistics_viz { + } + collective_device_list { + } + backend_config_payload { + id: 0 + } + } + program_shape { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + id: 1 + root_id: 4294967300 + } + host_program_shape { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + entry_computation_id: 1 + schedule { + sequences { + key: 0 + value { + instruction_ids: 1 + instruction_ids: 0 + instruction_ids: 2 + } + } + sequences { + key: 1 + value { + instruction_ids: 4294967296 + instruction_ids: 4294967298 + instruction_ids: 4294967297 + instruction_ids: 4294967299 + instruction_ids: 4294967300 + } + } + } + input_output_alias { + } + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + buffer_donor { + } + frontend_attributes { + map { + key: "fingerprint_before_lhs" + value: "c36d8870be58d238935bebefbbc5b827" + } + } + payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" + } + config { + entry_computation_layout { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + replica_count: 2 + num_partitions: 1 + intra_op_parallelism_threads: -1 + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.01 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "NEON" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: false + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: false + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + } + static_device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + allow_spmd_sharding_propagation_to_parameters: false + sharding_config { + } + memory_fitting_level: EFFORT_O2 + schedule_config { + } + partition_size: 2 + } + } + buffer_assignment { + logical_buffers { + size: 4 + defined_at { + } + } + logical_buffers { + id: 1 + size: 4 + defined_at { + instruction_id: 1 + } + } + logical_buffers { + id: 2 + size: 4 + defined_at { + instruction_id: 2 + } + } + logical_buffers { + id: 10 + size: 32 + defined_at { + instruction_id: 4294967297 + } + } + logical_buffers { + id: 11 + size: 32 + defined_at { + instruction_id: 4294967298 + } + } + logical_buffers { + id: 12 + size: 4 + defined_at { + instruction_id: 4294967296 + } + } + logical_buffers { + id: 13 + size: 32 + defined_at { + instruction_id: 4294967299 + } + } + logical_buffers { + id: 14 + size: 16 + defined_at { + instruction_id: 4294967300 + } + } + buffer_allocations { + size: 32 + assigned { + logical_buffer_id: 10 + size: 32 + element_type: U32 + } + is_constant: true + } + buffer_allocations { + index: 1 + size: 32 + assigned { + logical_buffer_id: 11 + size: 32 + element_type: U32 + } + is_constant: true + } + buffer_allocations { + index: 2 + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 12 + size: 4 + element_type: U32 + } + assigned { + logical_buffer_id: 14 + size: 16 + element_type: U32 + } + } + buffer_allocations { + index: 3 + size: 4 + is_thread_local: true + assigned { + logical_buffer_id: 2 + size: 4 + element_type: U32 + } + } + buffer_allocations { + index: 4 + size: 4 + is_thread_local: true + assigned { + size: 4 + element_type: U32 + } + } + buffer_allocations { + index: 5 + size: 4 + is_thread_local: true + assigned { + logical_buffer_id: 1 + size: 4 + element_type: U32 + } + } + buffer_allocations { + index: 6 + size: 32 + assigned { + logical_buffer_id: 13 + size: 32 + element_type: U32 + } + } + heap_simulator_traces { + events { + buffer_id: 13 + computation_name: "main" + instruction_name: "loop_select_fusion" + } + events { + kind: FREE + buffer_id: 13 + computation_name: "main" + instruction_name: "loop_select_fusion" + } + whole_module_simulation: true + buffer_allocation_index: 3 + } + peak_buffers { + logical_buffer_ids: 10 + } + peak_buffers { + index: 1 + logical_buffer_ids: 11 + } + peak_buffers { + index: 2 + logical_buffer_ids: 14 + } + peak_buffers { + index: 3 + logical_buffer_ids: 2 + } + peak_buffers { + index: 4 + logical_buffer_ids: 0 + } + peak_buffers { + index: 5 + logical_buffer_ids: 1 + } + peak_buffers { + index: 6 + logical_buffer_ids: 13 + } + } + binary: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000p\013\000\000\000\000\000\000\360\006\000\000\000\000\000\000\nd\000\006@\0008\000\004\000@\000\022\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.nv.shared.reserved.0\000.nv.global\000.nv.constant4\000\000\000\000.nv.compat\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant4\000\000\000\000.nv.callgraph\000.nv.prototype\000.nv.merc.debug_frame\000.nv.merc.rela.nv.constant.pic\000.nv.merc.nv.constant.pic\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.global\000.nv.constant.pic\000.nv.compat\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant.pic\000.nv.callgraph\000.nv.prototype\000buffer_for_c1_1\000buffer_for_c0_1\000.nv.constant4\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\200\000\000\000 \240\013\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000<\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\034\001\000\000\021\000\014\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000,\001\000\000\021\000\014\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\010\000\000\000\000\000\000\000\002\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\002\000\001\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\200\000\000\000 \240\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\254\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\034\001\000\000\035 \014\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000,\001\000\000\035 \014\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\\\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\001\000\000\000\000\000\000J\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000(\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\216\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\350\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\350\003\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\234\004\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\007\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\203\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\004\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\306\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\260\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\004\000\000\000\000\000\0000\000\000\000\000\000\000\000\003\000\000\000\n\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000r\000\000\000\001\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \005\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0000\005\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000g\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\000\000\000@\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\342\000\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\350\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\367\000\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\000\000\0000\000\000\000\000\000\000\000\021\000\000\000\017\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\001\000\000}\000\000p\002\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000 \005\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000L\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\006\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000p\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000p\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000 \005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\0000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\000\000\000\000\000\000\000\020\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + gpu_compute_capability { + cuda_compute_capability { + major: 10 + feature_extension: ACCELERATED_FEATURES + } + } + module_name: "test" + program_shape { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + buffer_allocations { + values { + size: 32 + is_constant: true + } + values { + index: 1 + size: 32 + is_constant: true + } + values { + index: 2 + size: 16 + maybe_live_out: true + } + values { + index: 3 + size: 4 + is_thread_local: true + } + values { + index: 4 + size: 4 + is_thread_local: true + } + values { + index: 5 + size: 4 + is_thread_local: true + } + values { + index: 6 + size: 32 + } + } + output_info_map { + shape_index { + } + output_info { + allocation_index: 2 + } + } + constants { + symbol_name: "buffer_for_c1_1" + allocation_index: 1 + } + constants { + symbol_name: "buffer_for_c0_1" + } + thunks { + thunk_info { + profile_annotation: "id.1" + thunk_id: 1 + } + replica_id_thunk { + dest_buffer { + size: 4 + buffer_allocation_index: 2 + element_type: U32 + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_select_fusion" + thunk_id: 2 + } + custom_kernel_thunk { + args { + slice { + size: 32 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + args { + slice { + size: 32 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + args { + slice { + size: 4 + buffer_allocation_index: 2 + element_type: U32 + } + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + } + args { + slice { + size: 32 + buffer_allocation_index: 6 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + written: false + written: false + written: false + written: true + custom_kernel { + name: "loop_select_fusion" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\030\025\000\000\000\000\000\000\230\017\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_select_fusion\000.nv.capmerc.text.loop_select_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_select_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_select_fusion\000.nv.constant0.loop_select_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\244\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\323\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000/\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000K\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000^\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\004<\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\014\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\360\000\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000\003\031 \000\004\n\010\000\007\000\000\000\200\003 \000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\344\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000.\016\000\202{\004\377\000\340\000\000\000\n\000\000\000b\016\000\031y\t\000\000\000\000\000\000!\000\000\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\314\017\000\202{\006\377\000\346\000\000\000\n\000\000\000\342\016\000\201y\002\002\004\000\000\000\000\231\036\014\000$\037\000\014r\000\002\377\000\000\000pP\360\003\000\332\017\001\202\013\004\377\000\342\000\000\000\010\000\000\000\260.\000\202\013\005\377\000\343\000\000\000\010\000\000\000\244\016\000%x\004\t\004\000\000\000\004\000\216\007\000\314O\000\201y\005\004\004\000\000\000\000\231\036\014\000\242\016\000%x\006\t\004\000\000\000\006\000\216\007\000\312\217\000\206y\000\006\005\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\021\000\000@\200\271\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\004\n\370\000\004\000\000\000A\002\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\203\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\030\000\000\000\002\"\016\006\000\000B\000\000\000\001\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\002\"\016\006\000\000B\000\000\000A\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\0028\0162\370\000@\021\000\000\000\000\202\001\n\000\000\002\001@\001\000\000\000\000\000\000\000\000\000\000\000P\010\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\001\000\000\000\000\000\000\004\020\000\000\000\004\220\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\014\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\220\001\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\244\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\323\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\000\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000/\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000K\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\240\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\350\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000L\002\000\000\000\000\000\000\177\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\305\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\370\004\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\254\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000j\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\364\005\000\000\000\000\000\000\204\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\272\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\365\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\343\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\000\002\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\244\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\021\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000\240\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\240\014\000\000\000\000\000\000\016\001\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000V\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\256\r\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000k\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000 \016\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000|\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000D\016\000\000\000\000\000\000\224\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\330\016\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\272\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\360\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\330\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\360\016\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\030\025\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\030\025\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\000\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\003\000\000\000\000\000\000\240\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 4 + kernel_name: "loop_select_fusion" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 2 + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 3 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 8 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "ars.3" + thunk_id: 3 + } + reduce_scatter_thunk { + collective_config { + operand_element_type: U32 + } + buffers { + element_count: 8 + source_buffer { + slice { + size: 32 + buffer_allocation_index: 6 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + destination_buffer { + slice { + size: 16 + buffer_allocation_index: 2 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + reduction_kind: REDUCTION_KIND_SUM + } + } + executable_abi_version { + platform_name: "CUDA" + cuda_platform_version { + cuda_toolkit_version: "12.8.0" + cudnn_version: "9.10.2" + cub_version: "3.1.2" + } + } + buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 64B\n maybe_live_out allocation: 16B\n preallocated temp allocation: 32B\n preallocated temp fragmentation: 0B (0.00%)\n total allocation: 124B\n total fragmentation: 76B (61.29%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 32B\n\t\tXLA Label: fusion\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: reduce-scatter\n\t\tShape: u32[4]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 7:\n\t\tSize: 4B\n\t\tXLA Label: add\n\t\tShape: u32[]\n\t\t==========================\n" +} diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v2/AllGatherMixedTypes.pbtxt b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v2/AllGatherMixedTypes.pbtxt new file mode 100644 index 00000000000000..4ea9837395c53d --- /dev/null +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v2/AllGatherMixedTypes.pbtxt @@ -0,0 +1,2190 @@ +# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto +# proto-message: xla.HumanReadableAotExecutable +executable_and_options { + compile_options { + executable_build_options { + device_ordinal: -1 + result_layout { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.1 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "NEON" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: true + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: true + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + xla_gpu_enable_persistent_symmetric_memory: false + xla_gpu_experimental_enable_raft_for_stable_topk: false + xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND + xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET + } + num_replicas: 2 + num_partitions: 1 + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + comp_envs { + } + allow_spmd_sharding_propagation_to_parameters: false + process_count: 1 + memory_fitting_level: EFFORT_O2 + } + } + pjrt_client_name: "PjRtStreamExecutorClient" +} +gpu_executable { + hlo_module_with_config { + hlo_module { + name: "test" + entry_computation_name: "test_computation" + computations { + name: "fused_transpose" + instructions { + name: "param_0.1" + opcode: "parameter" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_0.1" + } + id: 4294967298 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.6.1" + opcode: "bitcast" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.6.1" + } + id: 4294967297 + operand_ids: 4294967298 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "transpose.1.1" + opcode: "transpose" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "transpose.1.1" + } + dimensions: 1 + dimensions: 0 + id: 4294967296 + operand_ids: 4294967297 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + result { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.1" + } + id: 1 + root_id: 4294967296 + is_fusion_computation: true + } + computations { + name: "fused_transpose.1" + instructions { + name: "param_0.3" + opcode: "parameter" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_0.3" + } + id: 8589934594 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.4.1" + opcode: "bitcast" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.4.1" + } + id: 8589934593 + operand_ids: 8589934594 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "transpose.3" + opcode: "transpose" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "transpose.3" + } + dimensions: 1 + dimensions: 0 + id: 8589934592 + operand_ids: 8589934593 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + result { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.3" + } + id: 2 + root_id: 8589934592 + is_fusion_computation: true + } + computations { + name: "fused_broadcast" + instructions { + name: "param_0.5" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "param_0.5" + } + id: 12884901890 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "p1.1.1" + opcode: "convert" + shape { + element_type: F32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "p1.1.1" + } + id: 12884901889 + operand_ids: 12884901890 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "broadcast.13.1" + opcode: "broadcast" + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "broadcast.13.1" + } + id: 12884901888 + operand_ids: 12884901889 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.5" + } + id: 3 + root_id: 12884901888 + is_fusion_computation: true + } + computations { + name: "wrapped_broadcast_computation" + instructions { + name: "param_0.7" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "param_0.7" + } + id: 21474836481 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "broadcast.12.1" + opcode: "broadcast" + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "broadcast.12.1" + } + id: 21474836480 + operand_ids: 21474836481 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + parameter_names: "param_0.7" + } + id: 5 + root_id: 21474836480 + is_fusion_computation: true + } + computations { + name: "test_computation" + instructions { + name: "id.1" + opcode: "replica-id" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "id.1" + } + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "wrapped_broadcast" + opcode: "fusion" + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "wrapped_broadcast" + } + fusion_kind: "kLoop" + id: 9 + operand_ids: 0 + called_computation_ids: 5 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_broadcast_fusion" + opcode: "fusion" + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_broadcast_fusion" + } + fusion_kind: "kLoop" + id: 8 + operand_ids: 0 + called_computation_ids: 3 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "allgather.3" + opcode: "all-gather" + shape { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + metadata { + scheduling_name: "all-gather-start" + } + dimensions: 1 + id: 10 + operand_ids: 9 + operand_ids: 8 + frontend_attributes { + } + statistics_viz { + } + collective_device_list { + } + backend_config_payload { + id: 0 + } + } + instructions { + name: "ag0.1" + opcode: "get-tuple-element" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "ag0.1" + } + id: 1 + operand_ids: 10 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_transpose_fusion.1" + opcode: "fusion" + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_transpose_fusion.1" + } + fusion_kind: "kLoop" + id: 7 + operand_ids: 1 + called_computation_ids: 2 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.24" + opcode: "bitcast" + shape { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.24" + } + id: 4 + operand_ids: 7 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "ag1.1" + opcode: "get-tuple-element" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "ag1.1" + } + tuple_index: 1 + id: 2 + operand_ids: 10 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_transpose_fusion" + opcode: "fusion" + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_transpose_fusion" + } + fusion_kind: "kLoop" + id: 6 + operand_ids: 2 + called_computation_ids: 1 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "bitcast.1.0" + opcode: "bitcast" + shape { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "bitcast.1.0" + } + id: 5 + operand_ids: 6 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "out.1" + opcode: "tuple" + shape { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + metadata { + scheduling_name: "out.1" + } + id: 3 + operand_ids: 4 + operand_ids: 5 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + root_id: 3 + } + host_program_shape { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + schedule { + sequences { + key: 0 + value { + instruction_ids: 0 + instruction_ids: 8 + instruction_ids: 9 + instruction_ids: 10 + instruction_ids: 1 + instruction_ids: 7 + instruction_ids: 2 + instruction_ids: 6 + instruction_ids: 4 + instruction_ids: 5 + instruction_ids: 3 + } + } + } + input_output_alias { + } + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + buffer_donor { + } + frontend_attributes { + map { + key: "fingerprint_before_lhs" + value: "7b2e4f79c5524b5cbc566da178aa7fda" + } + } + payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" + } + config { + entry_computation_layout { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + replica_count: 2 + num_partitions: 1 + intra_op_parallelism_threads: -1 + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.1 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "NEON" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: true + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: true + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + xla_gpu_enable_persistent_symmetric_memory: false + xla_gpu_experimental_enable_raft_for_stable_topk: false + xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND + xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET + } + static_device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + allow_spmd_sharding_propagation_to_parameters: false + sharding_config { + } + memory_fitting_level: EFFORT_O2 + schedule_config { + } + partition_size: 2 + } + } + buffer_assignment { + logical_buffers { + id: 9 + size: 4 + defined_at { + } + } + logical_buffers { + id: 10 + size: 8 + defined_at { + instruction_id: 9 + } + } + logical_buffers { + id: 11 + size: 8 + defined_at { + instruction_id: 8 + } + } + logical_buffers { + id: 12 + size: 16 + defined_at { + instruction_id: 10 + } + } + logical_buffers { + id: 13 + size: 16 + defined_at { + shape_index: 0 + instruction_id: 10 + } + } + logical_buffers { + id: 14 + size: 16 + defined_at { + shape_index: 1 + instruction_id: 10 + } + } + logical_buffers { + id: 15 + size: 16 + defined_at { + instruction_id: 7 + } + } + logical_buffers { + id: 16 + size: 16 + defined_at { + instruction_id: 6 + } + } + logical_buffers { + id: 17 + size: 16 + defined_at { + instruction_id: 3 + } + } + buffer_allocations { + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 10 + size: 8 + element_type: U32 + } + assigned { + logical_buffer_id: 15 + size: 16 + element_type: U32 + } + } + buffer_allocations { + index: 1 + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 9 + size: 4 + element_type: U32 + } + assigned { + logical_buffer_id: 13 + size: 16 + element_type: U32 + } + assigned { + logical_buffer_id: 16 + size: 16 + element_type: F32 + } + } + buffer_allocations { + index: 2 + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 17 + size: 16 + element_type: TUPLE + } + is_tuple: true + } + buffer_allocations { + index: 3 + size: 520 + assigned { + logical_buffer_id: 11 + offset: 512 + size: 8 + element_type: F32 + } + assigned { + logical_buffer_id: 12 + size: 16 + element_type: TUPLE + } + assigned { + logical_buffer_id: 14 + offset: 256 + size: 16 + element_type: F32 + } + is_tuple: true + } + heap_simulator_traces { + events { + buffer_id: 11 + computation_name: "test_computation" + instruction_name: "loop_broadcast_fusion" + } + events { + buffer_id: 14 + computation_name: "test_computation" + instruction_name: "allgather.3" + } + events { + kind: FREE + buffer_id: 11 + computation_name: "test_computation" + instruction_name: "loop_broadcast_fusion" + } + events { + kind: FREE + buffer_id: 14 + computation_name: "test_computation" + instruction_name: "allgather.3" + } + whole_module_simulation: true + buffer_allocation_index: 3 + } + peak_buffers { + logical_buffer_ids: 15 + } + peak_buffers { + index: 1 + logical_buffer_ids: 16 + } + peak_buffers { + index: 2 + logical_buffer_ids: 17 + } + peak_buffers { + index: 3 + logical_buffer_ids: 12 + logical_buffer_ids: 11 + logical_buffer_ids: 14 + } + } + gpu_compute_capability { + cuda_compute_capability { + major: 10 + feature_extension: ACCELERATED_FEATURES + } + } + module_name: "test" + program_shape { + result { + element_type: TUPLE + tuple_shapes { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + tuple_shapes { + element_type: F32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + buffer_allocations { + values { + size: 16 + maybe_live_out: true + } + values { + index: 1 + size: 16 + maybe_live_out: true + } + values { + index: 2 + size: 16 + maybe_live_out: true + is_tuple: true + } + values { + index: 3 + size: 520 + is_tuple: true + } + } + output_info_map { + shape_index { + } + output_info { + allocation_index: 2 + } + } + output_info_map { + shape_index { + indexes: 0 + } + output_info { + } + } + output_info_map { + shape_index { + indexes: 1 + } + output_info { + allocation_index: 1 + } + } + thunks { + thunk_info { + profile_annotation: "id.1" + thunk_id: 1 + } + replica_id_thunk { + dest_buffer { + size: 4 + buffer_allocation_index: 1 + element_type: U32 + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_broadcast_fusion" + thunk_id: 2 + } + custom_kernel_thunk { + args { + slice { + size: 4 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + } + args { + slice { + offset: 512 + size: 8 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "loop_broadcast_fusion" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_broadcast_fusion\000.nv.capmerc.text.loop_broadcast_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_broadcast_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_broadcast_fusion\000.nv.constant0.loop_broadcast_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000j\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004(\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\240\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\340\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000\"\016\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\342\000\000\000\n\000\000\000b\016\000\201y\002\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\342/\000Er\007\000\002\000\000\000\000\020 \000\000\312O\000\206y\000\004\007\000\000\000\004\031\020\014\000\342\017\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\014\000\000@\300\005\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\001\000\000\000\000\000\000\000\000\000\000\000\320\006\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\004\020\000\000\000\004\360\000\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\360\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\372\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000^\002\000\000\000\000\000\000\216\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\316\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\030\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\314\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000m\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\006\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\303\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\376\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\354\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\032\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000e\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000z\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\213\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\262\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\314\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\352\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "loop_broadcast_fusion" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 2 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "wrapped_broadcast" + thunk_id: 3 + } + custom_kernel_thunk { + args { + slice { + size: 4 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + } + args { + slice { + size: 8 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "wrapped_broadcast" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.wrapped_broadcast\000.nv.capmerc.text.wrapped_broadcast\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.wrapped_broadcast\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000wrapped_broadcast\000.nv.constant0.wrapped_broadcast\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\320\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\374\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000Z\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004$\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\220\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\340\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000\"\016\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\342\000\000\000\n\000\000\000b\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\312/\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\013\000\000@\300\002\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\007\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\000\000\000\000\000\000\000\004\020\000\000\000\004\340\000\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\340\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\320\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\374\000\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\360\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\342\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000F\002\000\000\000\000\000\000z\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\302\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\350\004\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\234\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000i\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\344\005\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\267\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\362\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000X\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\340\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000f\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000w\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\232\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\264\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\322\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "wrapped_broadcast" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 2 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "allgather.3" + thunk_id: 4 + } + all_gather_thunk { + collective_config { + operand_element_type: U32 + operand_element_type: F32 + } + buffers { + element_count: 2 + source_buffer { + slice { + size: 8 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + destination_buffer { + slice { + size: 16 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + } + buffers { + element_count: 2 + source_buffer { + slice { + offset: 512 + size: 8 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 1 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + destination_buffer { + slice { + offset: 256 + size: 16 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_transpose_fusion.1" + thunk_id: 5 + } + custom_kernel_thunk { + args { + slice { + size: 16 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + args { + slice { + size: 16 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "loop_transpose_fusion_1" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion_1\000.nv.capmerc.text.loop_transpose_fusion_1\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_transpose_fusion_1\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion_1\000.nv.constant0.loop_transpose_fusion_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\342\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Z\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000r\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\340\000\000\000\n\000\000\000b\016\000\254w\004\377\000k\000\000\000\n\000\010\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\342\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242N\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\017\000\000@\340-\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\005\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\004\020\000\000\000\004`\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000`\001\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\342\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Z\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\006\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000j\002\000\000\000\000\000\000\230\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\004\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\324\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\0000\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\344\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\006\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000o\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\006\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\311\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\001\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\362\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000F\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000o\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\204\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\225\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\276\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\330\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "loop_transpose_fusion_1" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 4 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_transpose_fusion" + thunk_id: 6 + } + custom_kernel_thunk { + args { + slice { + offset: 256 + size: 16 + buffer_allocation_index: 3 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 0 + minor_to_major: 1 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + args { + slice { + size: 16 + buffer_allocation_index: 1 + element_type: F32 + } + shape { + element_type: F32 + dimensions: 2 + dimensions: 2 + layout { + minor_to_major: 1 + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + is_dynamic_dimension: false + } + } + written: false + written: true + custom_kernel { + name: "loop_transpose_fusion" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\010\024\000\000\000\000\000\000\210\016\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion\000.nv.capmerc.text.loop_transpose_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_transpose_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion\000.nv.constant0.loop_transpose_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000j\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\340\000\000\000\n\000\000\000b\016\000\254w\004\377\000k\000\000\000\n\000\010\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\342\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242N\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\017\000\000@\340-\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\005\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\004\020\000\000\000\004`\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000`\001\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\372\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000^\002\000\000\000\000\000\000\216\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\316\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\030\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\314\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000m\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\006\000\000\000\000\000\000d\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\303\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\376\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\354\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\032\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000>\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\256\000\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000e\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\276\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000z\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\213\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000T\r\000\000\000\000\000\000t\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\262\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\314\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\352\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\010\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 2 + kernel_name: "loop_transpose_fusion" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 4 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + executable_abi_version { + platform_name: "CUDA" + cuda_platform_version { + cuda_toolkit_version: "12.8.0" + cudnn_version: "9.10.2" + cub_version: "3.1.2" + } + } + buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 0B\n maybe_live_out allocation: 48B\n preallocated temp allocation: 520B\n preallocated temp fragmentation: 240B (46.15%)\n total allocation: 568B\n total fragmentation: 504B (88.73%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: (u32[2,2], f32[2,2])\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 16B\n\t\tXLA Label: tuple\n\t\tShape: (u32[4], f32[4])\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: u32[2,2]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 8B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,1]\n\t\t==========================\n" +} diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v2/ReduceScatter.pbtxt b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v2/ReduceScatter.pbtxt new file mode 100644 index 00000000000000..886fabba15a64e --- /dev/null +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_gb200/v2/ReduceScatter.pbtxt @@ -0,0 +1,1624 @@ +# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto +# proto-message: xla.HumanReadableAotExecutable +executable_and_options { + compile_options { + executable_build_options { + device_ordinal: -1 + result_layout { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.1 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "NEON" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: true + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: true + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + xla_gpu_enable_persistent_symmetric_memory: false + xla_gpu_experimental_enable_raft_for_stable_topk: false + xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND + xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET + } + num_replicas: 2 + num_partitions: 1 + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + comp_envs { + } + allow_spmd_sharding_propagation_to_parameters: false + process_count: 1 + memory_fitting_level: EFFORT_O2 + } + } + pjrt_client_name: "PjRtStreamExecutorClient" +} +gpu_executable { + hlo_module_with_config { + hlo_module { + name: "test" + entry_computation_name: "main" + computations { + name: "add" + instructions { + name: "lhs" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "lhs" + } + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "rhs" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "rhs" + } + parameter_number: 1 + id: 1 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "add.1" + opcode: "add" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "add.1" + } + id: 2 + operand_ids: 0 + operand_ids: 1 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + parameter_names: "lhs" + parameter_names: "rhs" + } + root_id: 2 + } + computations { + name: "fused_select" + instructions { + name: "param_2.2" + opcode: "parameter" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "param_2.2" + } + parameter_number: 2 + id: 8589934597 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "zero_2" + opcode: "constant" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + } + literal { + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + u32s: 0 + } + id: 8589934598 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "p.2" + opcode: "compare" + shape { + element_type: PRED + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "p.2" + } + id: 8589934596 + operand_ids: 8589934597 + operand_ids: 8589934598 + comparison_direction: "EQ" + frontend_attributes { + } + statistics_viz { + } + comparison_order: "TOTAL" + } + instructions { + name: "pb.2" + opcode: "broadcast" + shape { + element_type: PRED + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "pb.2" + } + id: 8589934593 + operand_ids: 8589934596 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "param_0.1" + opcode: "parameter" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_0.1" + } + id: 8589934594 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "param_1.1" + opcode: "parameter" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "param_1.1" + } + parameter_number: 1 + id: 8589934595 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "data.2" + opcode: "select" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "data.2" + } + id: 8589934592 + operand_ids: 8589934593 + operand_ids: 8589934594 + operand_ids: 8589934595 + frontend_attributes { + } + statistics_viz { + } + } + program_shape { + parameters { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + parameters { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + parameters { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + result { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + parameter_names: "param_0.1" + parameter_names: "param_1.1" + parameter_names: "param_2.2" + } + id: 2 + root_id: 8589934592 + is_fusion_computation: true + } + computations { + name: "main" + instructions { + name: "c0_1" + opcode: "constant" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + } + literal { + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + u32s: 1 + u32s: 2 + u32s: 3 + u32s: 4 + u32s: 5 + u32s: 6 + u32s: 7 + u32s: 8 + } + id: 4294967297 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "c1_1" + opcode: "constant" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + } + literal { + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + u32s: 10 + u32s: 11 + u32s: 12 + u32s: 13 + u32s: 14 + u32s: 15 + u32s: 16 + u32s: 17 + } + id: 4294967298 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "id.1" + opcode: "replica-id" + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + metadata { + scheduling_name: "id.1" + } + id: 4294967296 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "loop_select_fusion" + opcode: "fusion" + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "loop_select_fusion" + } + fusion_kind: "kLoop" + id: 4294967299 + operand_ids: 4294967297 + operand_ids: 4294967298 + operand_ids: 4294967296 + called_computation_ids: 2 + frontend_attributes { + } + statistics_viz { + } + } + instructions { + name: "ars.3" + opcode: "reduce-scatter" + shape { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + metadata { + scheduling_name: "reduce-scatter-start" + } + dimensions: 0 + id: 4294967300 + operand_ids: 4294967299 + called_computation_ids: 0 + frontend_attributes { + } + statistics_viz { + } + collective_device_list { + } + backend_config_payload { + id: 0 + } + } + program_shape { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + id: 1 + root_id: 4294967300 + } + host_program_shape { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + entry_computation_id: 1 + schedule { + sequences { + key: 0 + value { + instruction_ids: 1 + instruction_ids: 0 + instruction_ids: 2 + } + } + sequences { + key: 1 + value { + instruction_ids: 4294967296 + instruction_ids: 4294967298 + instruction_ids: 4294967297 + instruction_ids: 4294967299 + instruction_ids: 4294967300 + } + } + } + input_output_alias { + } + device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + buffer_donor { + } + frontend_attributes { + map { + key: "fingerprint_before_lhs" + value: "c36d8870be58d238935bebefbbc5b827" + } + } + payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" + } + config { + entry_computation_layout { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + replica_count: 2 + num_partitions: 1 + intra_op_parallelism_threads: -1 + debug_options { + xla_disable_hlo_passes: "constant_folding" + xla_backend_optimization_level: 3 + xla_eliminate_hlo_implicit_broadcast: true + xla_cpu_multi_thread_eigen: true + xla_gpu_cuda_data_dir: "./cuda_sdk_lib" + xla_llvm_enable_alias_scope_metadata: true + xla_llvm_enable_noalias_metadata: true + xla_llvm_enable_invariant_load_metadata: true + xla_llvm_disable_expensive_passes: false + xla_cpu_use_onednn: false + xla_cpu_enable_fast_math: false + xla_gpu_enable_fast_min_max: false + xla_force_host_platform_device_count: 1 + xla_gpu_enable_dynamic_slice_fusion: true + xla_hlo_evaluator_use_fast_path: true + xla_dump_hlo_as_html: false + xla_cpu_fast_math_honor_nans: true + xla_cpu_fast_math_honor_infs: true + xla_allow_excess_precision: true + xla_gpu_autotune_level: 4 + xla_cpu_fast_math_honor_division: true + xla_cpu_fast_math_honor_functions: true + xla_dump_include_timestamp: false + xla_dump_max_hlo_modules: -1 + xla_cpu_enable_xprof_traceme: false + xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false + xla_cpu_enable_fast_min_max: true + xla_multiheap_size_constraint_per_heap: -1 + xla_dump_module_metadata: false + xla_gpu_force_compilation_parallelism: 1 + xla_dump_fusion_visualization: false + xla_gpu_all_reduce_combine_threshold_bytes: 31457287 + xla_gpu_nccl_termination_timeout_seconds: -1 + xla_dump_hlo_as_long_text: true + xla_gpu_enable_cublaslt: true + xla_gpu_shape_checks: RUNTIME + xla_gpu_use_runtime_fusion: false + xla_dump_latency_hiding_schedule: false + xla_dump_enable_mlir_pretty_form: true + xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP + xla_gpu_enable_triton_gemm: true + xla_gpu_enable_cudnn_int8x32_convolution_reordering: true + xla_gpu_triton_gemm_any: true + xla_gpu_enable_while_loop_reduce_scatter_code_motion: false + xla_gpu_collective_inflation_factor: 1 + xla_gpu_graph_min_graph_size: 5 + xla_gpu_enable_reassociation_for_converted_ar: true + xla_gpu_pgle_profile_file_or_directory_path: "" + xla_gpu_all_gather_combine_threshold_bytes: 31457287 + xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 + xla_gpu_enable_highest_priority_async_stream: true + xla_gpu_exhaustive_tiling_search: false + xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 + xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 + xla_gpu_redzone_padding_bytes: 8388608 + xla_gpu_copy_insertion_use_region_analysis: false + xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 + xla_gpu_collect_cost_model_stats: false + xla_gpu_enable_pipelined_p2p: false + xla_gpu_cublas_fallback: true + xla_gpu_enable_while_loop_double_buffering: false + xla_gpu_filter_kernels_spilling_registers_on_autotuning: true + xla_debug_buffer_assignment_show_max: 15 + xla_detailed_logging: true + xla_enable_dumping: true + xla_gpu_enable_all_gather_combine_by_dim: false + xla_gpu_enable_analytical_latency_estimator: false + xla_gpu_llvm_verification_level: 0 + xla_gpu_enable_reduce_scatter_combine_by_dim: false + xla_gpu_enable_command_buffer: CONDITIONAL + xla_gpu_enable_command_buffer: CUBLAS + xla_gpu_enable_command_buffer: CUBLASLT + xla_gpu_enable_command_buffer: CUDNN + xla_gpu_enable_command_buffer: CUSTOM_CALL + xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION + xla_gpu_enable_command_buffer: FUSION + xla_gpu_enable_cub_radix_sort: true + xla_gpu_memory_limit_slop_factor: 95 + xla_gpu_target_config_filename: "" + xla_gpu_enable_cudnn_layer_norm: false + xla_gpu_threshold_for_windowed_einsum_mib: 100000 + xla_gpu_enable_nccl_user_buffers: false + xla_gpu_enable_libnvptxcompiler: true + xla_gpu_enable_nccl_comm_splitting: true + xla_gpu_nccl_collective_max_nchannels: 0 + xla_gpu_nccl_p2p_max_nchannels: 0 + xla_gpu_nccl_init_max_rank_per_root_ratio: 0 + xla_gpu_multi_streamed_windowed_einsum: true + xla_gpu_gemm_rewrite_size_threshold: 100 + xla_gpu_require_complete_aot_autotune_results: false + xla_gpu_cudnn_gemm_fusion_level: 0 + xla_gpu_use_memcpy_local_p2p: false + xla_gpu_autotune_max_solutions: 0 + xla_dump_large_constants: false + xla_gpu_verify_triton_fusion_numerics: false + xla_reduce_window_rewrite_base_length: 16 + xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL + xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true + xla_gpu_nccl_terminate_on_error: false + xla_gpu_shard_autotuning: true + xla_gpu_enable_approx_costly_collectives: false + xla_cpu_prefer_vector_width: 256 + xla_gpu_per_fusion_autotune_cache_dir: "" + xla_cmd_buffer_trace_cache_size: 16 + xla_gpu_temp_buffer_use_separate_color: false + xla_syntax_sugar_async_ops: false + xla_gpu_autotune_gemm_rtol: 0.1 + xla_enable_command_buffers_during_profiling: false + xla_gpu_cudnn_gemm_max_plans: 5 + xla_cpu_parallel_codegen_split_count: 32 + xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE + xla_gpu_executable_warn_stuck_timeout_seconds: 10 + xla_gpu_executable_terminate_timeout_seconds: 30 + xla_gpu_experimental_disable_binary_libraries: false + xla_gpu_dot_merger_threshold_mb: 64 + xla_cpu_max_isa: "NEON" + xla_gpu_experimental_enable_fusion_block_level_rewriter: false + xla_enable_fast_math: false + xla_gpu_experimental_parallel_collective_overlap_limit: 1 + xla_cpu_copy_insertion_use_region_analysis: false + xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1 + xla_gpu_experimental_enable_triton_heroless_priority_fusion: false + xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN + xla_gpu_experimental_stream_annotation: true + xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO + xla_pjrt_allow_auto_layout_in_hlo: false + xla_gpu_enable_scatter_determinism_expander: false + xla_gpu_require_exclusive_lock: false + xla_gpu_generate_debug_info: false + xla_gpu_generate_line_info: false + xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false + xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE + xla_gpu_fail_ptx_compilation_on_register_spilling: false + xla_gpu_enable_analytical_sol_latency_estimator: true + xla_gpu_analytical_latency_estimator_options { + key: "chunk_prep_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "chunk_size_bytes" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "gpus_per_node" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nccl_op_launch_us" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "nic_speed_gbps" + value: "-1" + } + xla_gpu_analytical_latency_estimator_options { + key: "rtt_us" + value: "-1" + } + xla_cpu_use_xnnpack: true + xla_gpu_experimental_pack_dot_operands_along_k_dimension: true + xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false + xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED + xla_gpu_experimental_enable_heuristic_collective_combining: false + xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false + xla_hlo_pass_fix_detect_cycles: false + xla_cpu_generate_unique_c_style_kernel_entry_points: false + xla_gpu_experimental_collective_cse_distance_threshold: 0 + xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true + xla_gpu_experimental_collective_perf_table_path: "" + xla_gpu_collective_permute_combine_threshold_bytes: 31457287 + xla_unsupported_crash_on_hlo_pass_noop_change: false + xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false + xla_dump_full_hlo_config: true + xla_gpu_unsupported_enable_triton_multi_output_fusion: false + xla_gpu_experimental_matmul_perf_table_path: "" + xla_gpu_unsupported_enable_all_reduce_decomposer: false + xla_gpu_nccl_blocking_communicators: true + xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 + xla_gpu_first_collective_call_terminate_timeout_seconds: 40 + xla_gpu_nccl_async_execution: false + xla_cpu_emitter_verification_level: 1 + xla_early_exit_with_layouts: false + xla_gpu_experimental_use_ragged_dot_fusion: false + xla_detect_unstable_reductions: DETECTION_MODE_NONE + xla_gpu_command_buffer_scheduling_mode: LHS + xla_gpu_experimental_enable_nccl_symmetric_buffers: false + xla_disable_automatic_host_compute_offload: false + xla_gpu_experimental_enable_fusion_autotuner: true + xla_gpu_experimental_scaled_dot_with_triton: true + xla_gpu_command_buffer_unroll_loops: false + xla_cpu_experimental_onednn_custom_call: false + xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true + xla_cpu_collective_call_terminate_timeout_seconds: 40 + xla_cpu_collective_call_warn_stuck_seconds: 20 + xla_keep_shardings_after_spmd: false + xla_gpu_experimental_enable_triton_warp_specialization: false + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION + xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE + xla_cpu_enable_platform_dependent_math: true + xla_gpu_detect_nan: DETECTION_MODE_NONE + xla_gpu_detect_inf: DETECTION_MODE_NONE + xla_enable_enzyme_comms_opt: false + xla_gpu_experimental_enable_buffer_saver_on_thunks: false + xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE + xla_enable_scoped_logging_timers: true + xla_gpu_executable_embed_debug_info: true + xla_cpu_collective_timeout_seconds: 1800 + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false + xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU + xla_gpu_experimental_all_fusions_with_triton: false + xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT + xla_gpu_rocm_max_trace_events: 4194304 + xla_gpu_execution_terminate_timeout: "inf" + xla_enable_hlo_sharding_v3: true + xla_gpu_log_minmax: false + xla_gpu_print_compilation_stats: false + xla_gpu_experimental_enable_conv_fusion: false + xla_gpu_experimental_enable_tiling_propagation: true + xla_gpu_execution_progress_tracking: 0 + xla_gpu_experimental_max_unroll_factor: 32 + xla_gpu_enable_pdl: true + xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true + xla_recognize_reduction_optimization_level: 0 + xla_gpu_experimental_parallel_async_compute_limit: 2 + xla_dump_buffer_assignment_analysis: true + xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE + xla_gpu_experimental_enable_collective_multi_streaming: false + xla_gpu_collective_permute_connected_components: false + xla_gpu_experimental_force_split_k: 0 + xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_experimental_gemm_fusion_v2: false + xla_gpu_fusion_autotune_top_k_configs: 1 + xla_gpu_native_emitter_tune_unroll_factor_for_loops: false + xla_enable_rgv3_materialization: true + xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK + xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false + xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY + xla_gpu_executable_num_compute_streams: 0 + xla_gpu_executable_num_communication_streams: 0 + xla_gpu_experimental_collective_start_as_early_as_possible: false + xla_cpu_experimental_enable_tiling_propagation: true + xla_gpu_enable_dus_accumulator_zero_init_elimination: false + xla_gpu_experimental_thunk_buffer_debug_module_outputs: false + xla_cpu_ftz: true + xla_gpu_enable_allocator_spatial_partitioning: true + xla_dump_compact_gte: true + xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES + xla_dump_hlo_as_riegeli: false + xla_gpu_experimental_use_ragged_dot_grouped_gemm: true + xla_dump_hlo_to_subfolder: false + xla_gpu_enable_pdl_launch: true + xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED + xla_gpu_enable_nccl_user_buffers_in_default_space: false + xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + xla_gpu_enable_persistent_symmetric_memory: false + xla_gpu_experimental_enable_raft_for_stable_topk: false + xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND + xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET + } + static_device_assignment { + replica_count: 2 + computation_count: 1 + computation_devices { + replica_device_ids: 0 + replica_device_ids: 1 + } + } + allow_spmd_sharding_propagation_to_output: false + allow_spmd_sharding_propagation_to_parameters: false + sharding_config { + } + memory_fitting_level: EFFORT_O2 + schedule_config { + } + partition_size: 2 + } + } + buffer_assignment { + logical_buffers { + size: 4 + defined_at { + } + } + logical_buffers { + id: 1 + size: 4 + defined_at { + instruction_id: 1 + } + } + logical_buffers { + id: 2 + size: 4 + defined_at { + instruction_id: 2 + } + } + logical_buffers { + id: 10 + size: 32 + defined_at { + instruction_id: 4294967297 + } + } + logical_buffers { + id: 11 + size: 32 + defined_at { + instruction_id: 4294967298 + } + } + logical_buffers { + id: 12 + size: 4 + defined_at { + instruction_id: 4294967296 + } + } + logical_buffers { + id: 13 + size: 32 + defined_at { + instruction_id: 4294967299 + } + } + logical_buffers { + id: 14 + size: 16 + defined_at { + instruction_id: 4294967300 + } + } + buffer_allocations { + size: 32 + assigned { + logical_buffer_id: 10 + size: 32 + element_type: U32 + } + is_constant: true + } + buffer_allocations { + index: 1 + size: 32 + assigned { + logical_buffer_id: 11 + size: 32 + element_type: U32 + } + is_constant: true + } + buffer_allocations { + index: 2 + size: 16 + maybe_live_out: true + assigned { + logical_buffer_id: 12 + size: 4 + element_type: U32 + } + assigned { + logical_buffer_id: 14 + size: 16 + element_type: U32 + } + } + buffer_allocations { + index: 3 + size: 4 + is_thread_local: true + assigned { + logical_buffer_id: 2 + size: 4 + element_type: U32 + } + } + buffer_allocations { + index: 4 + size: 4 + is_thread_local: true + assigned { + size: 4 + element_type: U32 + } + } + buffer_allocations { + index: 5 + size: 4 + is_thread_local: true + assigned { + logical_buffer_id: 1 + size: 4 + element_type: U32 + } + } + buffer_allocations { + index: 6 + size: 32 + assigned { + logical_buffer_id: 13 + size: 32 + element_type: U32 + } + } + heap_simulator_traces { + events { + buffer_id: 13 + computation_name: "main" + instruction_name: "loop_select_fusion" + } + events { + kind: FREE + buffer_id: 13 + computation_name: "main" + instruction_name: "loop_select_fusion" + } + whole_module_simulation: true + buffer_allocation_index: 3 + } + peak_buffers { + logical_buffer_ids: 10 + } + peak_buffers { + index: 1 + logical_buffer_ids: 11 + } + peak_buffers { + index: 2 + logical_buffer_ids: 14 + } + peak_buffers { + index: 3 + logical_buffer_ids: 2 + } + peak_buffers { + index: 4 + logical_buffer_ids: 0 + } + peak_buffers { + index: 5 + logical_buffer_ids: 1 + } + peak_buffers { + index: 6 + logical_buffer_ids: 13 + } + } + binary: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000p\013\000\000\000\000\000\000\360\006\000\000\000\000\000\000\nd\000\006@\0008\000\004\000@\000\022\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.nv.shared.reserved.0\000.nv.global\000.nv.constant4\000\000\000\000.nv.compat\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant4\000\000\000\000.nv.callgraph\000.nv.prototype\000.nv.merc.debug_frame\000.nv.merc.rela.nv.constant.pic\000.nv.merc.nv.constant.pic\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.global\000.nv.constant.pic\000.nv.compat\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant.pic\000.nv.callgraph\000.nv.prototype\000buffer_for_c1_1\000buffer_for_c0_1\000.nv.constant4\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\200\000\000\000 \240\013\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000<\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\034\001\000\000\021\000\014\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000,\001\000\000\021\000\014\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\010\000\000\000\000\000\000\000\002\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\002\000\001\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\200\000\000\000 \240\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\254\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\034\001\000\000\035 \014\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000,\001\000\000\035 \014\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\\\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\001\000\000\000\000\000\000J\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000(\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\216\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\350\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\350\003\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\234\004\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\007\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\203\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\004\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\306\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\260\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\004\000\000\000\000\000\0000\000\000\000\000\000\000\000\003\000\000\000\n\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000r\000\000\000\001\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \005\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0000\005\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000g\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\000\000\000@\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\342\000\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\350\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\367\000\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\000\000\0000\000\000\000\000\000\000\000\021\000\000\000\017\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\001\000\000}\000\000p\002\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000 \005\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000L\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0000\006\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000p\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000p\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000 \005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\0000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\000\000\000\000\000\000\000\020\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + gpu_compute_capability { + cuda_compute_capability { + major: 10 + feature_extension: ACCELERATED_FEATURES + } + } + module_name: "test" + program_shape { + result { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + buffer_allocations { + values { + size: 32 + is_constant: true + } + values { + index: 1 + size: 32 + is_constant: true + } + values { + index: 2 + size: 16 + maybe_live_out: true + } + values { + index: 3 + size: 4 + is_thread_local: true + } + values { + index: 4 + size: 4 + is_thread_local: true + } + values { + index: 5 + size: 4 + is_thread_local: true + } + values { + index: 6 + size: 32 + } + } + output_info_map { + shape_index { + } + output_info { + allocation_index: 2 + } + } + constants { + symbol_name: "buffer_for_c1_1" + allocation_index: 1 + } + constants { + symbol_name: "buffer_for_c0_1" + } + thunks { + thunk_info { + profile_annotation: "id.1" + thunk_id: 1 + } + replica_id_thunk { + dest_buffer { + size: 4 + buffer_allocation_index: 2 + element_type: U32 + } + } + } + thunks { + thunk_info { + profile_annotation: "loop_select_fusion" + thunk_id: 2 + } + custom_kernel_thunk { + args { + slice { + size: 32 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + args { + slice { + size: 32 + buffer_allocation_index: 1 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + args { + slice { + size: 4 + buffer_allocation_index: 2 + element_type: U32 + } + shape { + element_type: U32 + layout { + tail_padding_alignment_in_elements: 1 + } + } + } + args { + slice { + size: 32 + buffer_allocation_index: 6 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + written: false + written: false + written: false + written: true + custom_kernel { + name: "loop_select_fusion" + kernel_spec { + cubin { + data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\030\025\000\000\000\000\000\000\230\017\000\000\000\000\000\000\nd\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.shared.reserved.0\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_select_fusion\000.nv.capmerc.text.loop_select_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_select_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuver\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.compat\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_select_fusion\000.nv.constant0.loop_select_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\244\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\323\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000/\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000K\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000^\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\004<\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\200\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 12.8, V12.8.93\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills -arch sm_100a \000\000\014\000\000\000\014\000\000\000\350\003\000\000NVIDIA Corp\000\001\000d\000\001\000\000\000\001\000\000\000\004/\010\000\006\000\000\000\014\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\360\000\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000\003\031 \000\004\n\010\000\007\000\000\000\200\003 \000\0046\004\000\010\000\000\000\002\002\001\000\002\005\005\000\002\003\000\000\002\006\001\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\344\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000.\016\000\202{\004\377\000\340\000\000\000\n\000\000\000b\016\000\031y\t\000\000\000\000\000\000!\000\000\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\314\017\000\202{\006\377\000\346\000\000\000\n\000\000\000\342\016\000\201y\002\002\004\000\000\000\000\231\036\014\000$\037\000\014r\000\002\377\000\000\000pP\360\003\000\332\017\001\202\013\004\377\000\342\000\000\000\010\000\000\000\260.\000\202\013\005\377\000\343\000\000\000\010\000\000\000\244\016\000%x\004\t\004\000\000\000\004\000\216\007\000\314O\000\201y\005\004\004\000\000\000\000\231\036\014\000\242\016\000%x\006\t\004\000\000\000\006\000\216\007\000\312\217\000\206y\000\006\005\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\021\000\000@\200\271\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\004\n\370\000\004\000\000\000A\002\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\203\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\030\000\000\000\002\"\016\006\000\000B\000\000\000\001\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\002\"\016\006\000\000B\000\000\000A\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\0028\0162\370\000@\021\000\000\000\000\202\001\n\000\000\002\001@\001\000\000\000\000\000\000\000\000\000\000\000P\010\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\001\000\000\000\000\000\000\004\020\000\000\000\004\220\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\006\000\000\000\014\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004Z$\000\212\235\"\244\261\235\024m\000\264*\363\367X\003\216\014\007\n\033\342\336\212\327Rc\207\014\327+\007\000\315+\000\000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\002J\000\000\004\034\004\000\220\001\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000D\000\000\000\000\000\000\000=\000\001\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000Q\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\244\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\323\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\000\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000/\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000K\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\240\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\350\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000L\002\000\000\000\000\000\000\177\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\305\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\370\004\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\254\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\005\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000j\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\364\005\000\000\000\000\000\000\204\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\272\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\365\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\343\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000Q\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\000\002\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\244\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\021\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000\240\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\240\014\000\000\000\000\000\000\016\001\000\000\000\000\000\000\025\000\000\000\006\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000V\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\256\r\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000k\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000 \016\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000|\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000D\016\000\000\000\000\000\000\224\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\330\016\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\272\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\360\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\330\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\360\016\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\030\025\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\030\025\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\000\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\003\000\000\000\000\000\000\240\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + } + arity: 4 + kernel_name: "loop_select_fusion" + kernel_args_packing_spec { + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 1 + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 2 + } + } + kernel_arguments { + relocations { + kind: KIND_BITS64_ABSOLUTE + argument_index: 3 + } + } + } + } + block_dims { + coordinates { + x: 1 + y: 1 + z: 1 + } + } + thread_dims { + coordinates { + x: 8 + y: 1 + z: 1 + } + } + } + use_pdl: true + tma_metadata { + } + } + } + thunks { + thunk_info { + profile_annotation: "ars.3" + thunk_id: 3 + } + reduce_scatter_thunk { + collective_config { + operand_element_type: U32 + } + buffers { + element_count: 8 + source_buffer { + slice { + size: 32 + buffer_allocation_index: 6 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 8 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + destination_buffer { + slice { + size: 16 + buffer_allocation_index: 2 + element_type: U32 + } + shape { + element_type: U32 + dimensions: 4 + layout { + minor_to_major: 0 + tail_padding_alignment_in_elements: 1 + } + is_dynamic_dimension: false + } + } + } + reduction_kind: REDUCTION_KIND_SUM + } + } + executable_abi_version { + platform_name: "CUDA" + cuda_platform_version { + cuda_toolkit_version: "12.8.0" + cudnn_version: "9.10.2" + cub_version: "3.1.2" + } + } + buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 64B\n maybe_live_out allocation: 16B\n preallocated temp allocation: 32B\n preallocated temp fragmentation: 0B (0.00%)\n total allocation: 124B\n total fragmentation: 76B (61.29%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 32B\n\t\tXLA Label: fusion\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: reduce-scatter\n\t\tShape: u32[4]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 7:\n\t\tSize: 4B\n\t\tXLA Label: add\n\t\tShape: u32[]\n\t\t==========================\n" +} diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu/v2/AllGatherMixedTypes.pbtxt b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_h100/v1/AllGatherMixedTypes.pbtxt similarity index 60% rename from third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu/v2/AllGatherMixedTypes.pbtxt rename to third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_h100/v1/AllGatherMixedTypes.pbtxt index de78ec91592c00..cb7e12a41013aa 100644 --- a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu/v2/AllGatherMixedTypes.pbtxt +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_h100/v1/AllGatherMixedTypes.pbtxt @@ -1,9 +1,9 @@ -# proto-file: third_party/tensorflow/compiler/xla/tools/split_proto/split_proto_cli.proto -# proto-message: xla.split_proto_cli.DeserializedSplitExecutableAndOptions - +# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto +# proto-message: xla.HumanReadableAotExecutable executable_and_options { compile_options { executable_build_options { + device_ordinal: -1 result_layout { element_type: TUPLE tuple_shapes { @@ -39,7 +39,7 @@ executable_and_options { xla_cpu_enable_fast_math: false xla_gpu_enable_fast_min_max: false xla_force_host_platform_device_count: 1 - xla_gpu_enable_dynamic_slice_fusion: false + xla_gpu_enable_dynamic_slice_fusion: true xla_hlo_evaluator_use_fast_path: true xla_dump_hlo_as_html: false xla_cpu_fast_math_honor_nans: true @@ -65,7 +65,6 @@ executable_and_options { xla_gpu_use_runtime_fusion: false xla_dump_latency_hiding_schedule: false xla_dump_enable_mlir_pretty_form: true - xla_gpu_enable_latency_hiding_scheduler: false xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP xla_gpu_enable_triton_gemm: true xla_gpu_enable_cudnn_int8x32_convolution_reordering: true @@ -78,14 +77,10 @@ executable_and_options { xla_gpu_all_gather_combine_threshold_bytes: 31457287 xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 xla_gpu_enable_highest_priority_async_stream: true - xla_gpu_enable_pipelined_all_reduce: false xla_gpu_exhaustive_tiling_search: false xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 - xla_gpu_enable_pipelined_all_gather: false xla_gpu_redzone_padding_bytes: 8388608 - xla_gpu_enable_pipelined_reduce_scatter: true - xla_gpu_fused_attention_use_cudnn_rng: false xla_gpu_copy_insertion_use_region_analysis: false xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 xla_gpu_collect_cost_model_stats: false @@ -129,6 +124,7 @@ executable_and_options { xla_reduce_window_rewrite_base_length: 16 xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true xla_gpu_nccl_terminate_on_error: false xla_gpu_shard_autotuning: true xla_gpu_enable_approx_costly_collectives: false @@ -137,7 +133,7 @@ executable_and_options { xla_cmd_buffer_trace_cache_size: 16 xla_gpu_temp_buffer_use_separate_color: false xla_syntax_sugar_async_ops: false - xla_gpu_autotune_gemm_rtol: 0.1 + xla_gpu_autotune_gemm_rtol: 0.01 xla_enable_command_buffers_during_profiling: false xla_gpu_cudnn_gemm_max_plans: 5 xla_cpu_parallel_codegen_split_count: 32 @@ -199,7 +195,6 @@ executable_and_options { xla_cpu_generate_unique_c_style_kernel_entry_points: false xla_gpu_experimental_collective_cse_distance_threshold: 0 xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true - xla_cpu_use_fusion_emitters: true xla_gpu_experimental_collective_perf_table_path: "" xla_gpu_collective_permute_combine_threshold_bytes: 31457287 xla_unsupported_crash_on_hlo_pass_noop_change: false @@ -208,7 +203,6 @@ executable_and_options { xla_gpu_unsupported_enable_triton_multi_output_fusion: false xla_gpu_experimental_matmul_perf_table_path: "" xla_gpu_unsupported_enable_all_reduce_decomposer: false - xla_gpu_unsupported_use_all_reduce_one_shot_kernel: true xla_gpu_nccl_blocking_communicators: true xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 xla_gpu_first_collective_call_terminate_timeout_seconds: 40 @@ -219,14 +213,13 @@ executable_and_options { xla_detect_unstable_reductions: DETECTION_MODE_NONE xla_gpu_command_buffer_scheduling_mode: LHS xla_gpu_experimental_enable_nccl_symmetric_buffers: false - xla_gpu_experimental_autotuner_cache_dir: "" xla_disable_automatic_host_compute_offload: false xla_gpu_experimental_enable_fusion_autotuner: true xla_gpu_experimental_scaled_dot_with_triton: true xla_gpu_command_buffer_unroll_loops: false xla_cpu_experimental_onednn_custom_call: false - xla_gpu_experimental_use_raft_select_k: false xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true xla_cpu_collective_call_terminate_timeout_seconds: 40 xla_cpu_collective_call_warn_stuck_seconds: 20 xla_keep_shardings_after_spmd: false @@ -236,21 +229,32 @@ executable_and_options { xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE xla_cpu_enable_platform_dependent_math: true xla_gpu_detect_nan: DETECTION_MODE_NONE - xla_gpu_experimental_dump_gpu_executable: true xla_gpu_detect_inf: DETECTION_MODE_NONE xla_enable_enzyme_comms_opt: false xla_gpu_experimental_enable_buffer_saver_on_thunks: false xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE - xla_gpu_experimental_aot_compiled_thunks: true xla_enable_scoped_logging_timers: true xla_gpu_executable_embed_debug_info: true xla_cpu_collective_timeout_seconds: 1800 xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK xla_gpu_experimental_all_fusions_with_triton: false xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT xla_gpu_rocm_max_trace_events: 4194304 - xla_gpu_experimental_ragged_all_to_all_use_barrier: true xla_gpu_execution_terminate_timeout: "inf" xla_enable_hlo_sharding_v3: false xla_gpu_log_minmax: false @@ -261,6 +265,7 @@ executable_and_options { xla_gpu_experimental_max_unroll_factor: 32 xla_gpu_enable_pdl: true xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true xla_recognize_reduction_optimization_level: 0 xla_gpu_experimental_parallel_async_compute_limit: 2 xla_dump_buffer_assignment_analysis: true @@ -280,12 +285,12 @@ executable_and_options { xla_gpu_executable_num_compute_streams: 0 xla_gpu_executable_num_communication_streams: 0 xla_gpu_experimental_collective_start_as_early_as_possible: false - xla_cpu_experimental_enable_tiling_propagation: false + xla_cpu_experimental_enable_tiling_propagation: true xla_gpu_enable_dus_accumulator_zero_init_elimination: false xla_gpu_experimental_thunk_buffer_debug_module_outputs: false xla_cpu_ftz: true xla_gpu_enable_allocator_spatial_partitioning: true - xla_dump_compact_gte: false + xla_dump_compact_gte: true xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES xla_dump_hlo_as_riegeli: false xla_gpu_experimental_use_ragged_dot_grouped_gemm: true @@ -294,6 +299,35 @@ executable_and_options { xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED xla_gpu_enable_nccl_user_buffers_in_default_space: false xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" } num_replicas: 2 num_partitions: 1 @@ -311,9 +345,9 @@ executable_and_options { allow_spmd_sharding_propagation_to_parameters: false process_count: 1 memory_fitting_level: EFFORT_O2 - slice_size: 2 } } + pjrt_client_name: "PjRtStreamExecutorClient" } gpu_executable { hlo_module_with_config { @@ -338,6 +372,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "param_0.1" } id: 4294967298 frontend_attributes { @@ -361,6 +396,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "bitcast.6.1" } id: 4294967297 operand_ids: 4294967298 @@ -385,6 +421,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "transpose.1.1" } dimensions: 1 dimensions: 0 @@ -444,6 +481,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "param_0.3" } id: 8589934594 frontend_attributes { @@ -467,6 +505,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "bitcast.4.1" } id: 8589934593 operand_ids: 8589934594 @@ -491,6 +530,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "transpose.3" } dimensions: 1 dimensions: 0 @@ -544,6 +584,7 @@ gpu_executable { } } metadata { + scheduling_name: "param_0.5" } id: 12884901890 frontend_attributes { @@ -561,6 +602,7 @@ gpu_executable { } } metadata { + scheduling_name: "p1.1.1" } id: 12884901889 operand_ids: 12884901890 @@ -585,6 +627,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "broadcast.13.1" } id: 12884901888 operand_ids: 12884901889 @@ -618,156 +661,6 @@ gpu_executable { root_id: 12884901888 is_fusion_computation: true } - computations { - name: "async_computation" - instructions { - name: "param_0.6" - opcode: "parameter" - shape { - element_type: U32 - dimensions: 2 - dimensions: 1 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - metadata { - } - id: 17179869184 - frontend_attributes { - } - statistics_viz { - } - } - instructions { - name: "param_1" - opcode: "parameter" - shape { - element_type: F32 - dimensions: 2 - dimensions: 1 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - metadata { - } - parameter_number: 1 - id: 17179869185 - frontend_attributes { - } - statistics_viz { - } - } - instructions { - name: "allgather.2" - opcode: "all-gather" - shape { - element_type: TUPLE - tuple_shapes { - element_type: U32 - dimensions: 2 - dimensions: 2 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - tuple_shapes { - element_type: F32 - dimensions: 2 - dimensions: 2 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - } - metadata { - } - dimensions: 1 - id: 17179869186 - operand_ids: 17179869184 - operand_ids: 17179869185 - frontend_attributes { - } - statistics_viz { - } - collective_device_list { - } - } - program_shape { - parameters { - element_type: U32 - dimensions: 2 - dimensions: 1 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - parameters { - element_type: F32 - dimensions: 2 - dimensions: 1 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - result { - element_type: TUPLE - tuple_shapes { - element_type: U32 - dimensions: 2 - dimensions: 2 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - tuple_shapes { - element_type: F32 - dimensions: 2 - dimensions: 2 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - } - parameter_names: "param_0.6" - parameter_names: "param_1" - } - id: 4 - root_id: 17179869186 - } computations { name: "wrapped_broadcast_computation" instructions { @@ -780,6 +673,7 @@ gpu_executable { } } metadata { + scheduling_name: "param_0.7" } id: 21474836481 frontend_attributes { @@ -803,6 +697,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "broadcast.12.1" } id: 21474836480 operand_ids: 21474836481 @@ -848,6 +743,7 @@ gpu_executable { } } metadata { + scheduling_name: "id.1" } frontend_attributes { } @@ -870,18 +766,16 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "wrapped_broadcast" } fusion_kind: "kLoop" - id: 11 + id: 9 operand_ids: 0 called_computation_ids: 5 frontend_attributes { } statistics_viz { } - backend_config_payload { - id: 0 - } } instructions { name: "loop_broadcast_fusion" @@ -899,6 +793,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "loop_broadcast_fusion" } fusion_kind: "kLoop" id: 8 @@ -908,87 +803,10 @@ gpu_executable { } statistics_viz { } - backend_config_payload { - id: 0 - } - } - instructions { - name: "all-gather-start" - opcode: "async-start" - shape { - element_type: TUPLE - tuple_shapes { - element_type: TUPLE - tuple_shapes { - element_type: U32 - dimensions: 2 - dimensions: 1 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - tuple_shapes { - element_type: F32 - dimensions: 2 - dimensions: 1 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - } - tuple_shapes { - element_type: TUPLE - tuple_shapes { - element_type: U32 - dimensions: 2 - dimensions: 2 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - tuple_shapes { - element_type: F32 - dimensions: 2 - dimensions: 2 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - } - } - metadata { - } - id: 9 - operand_ids: 11 - operand_ids: 8 - called_computation_ids: 4 - frontend_attributes { - } - statistics_viz { - } - backend_config_payload { - id: 1 - } } instructions { - name: "all-gather-done" - opcode: "async-done" + name: "allgather.3" + opcode: "all-gather" shape { element_type: TUPLE tuple_shapes { @@ -1017,13 +835,21 @@ gpu_executable { } } metadata { + scheduling_name: "all-gather-start" } + dimensions: 1 id: 10 operand_ids: 9 + operand_ids: 8 frontend_attributes { } statistics_viz { } + collective_device_list { + } + backend_config_payload { + id: 0 + } } instructions { name: "ag0.1" @@ -1041,6 +867,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "ag0.1" } id: 1 operand_ids: 10 @@ -1065,6 +892,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "loop_transpose_fusion.1" } fusion_kind: "kLoop" id: 7 @@ -1074,9 +902,6 @@ gpu_executable { } statistics_viz { } - backend_config_payload { - id: 0 - } } instructions { name: "bitcast.24" @@ -1091,6 +916,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "bitcast.24" } id: 4 operand_ids: 7 @@ -1115,6 +941,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "ag1.1" } tuple_index: 1 id: 2 @@ -1140,6 +967,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "loop_transpose_fusion" } fusion_kind: "kLoop" id: 6 @@ -1149,9 +977,6 @@ gpu_executable { } statistics_viz { } - backend_config_payload { - id: 0 - } } instructions { name: "bitcast.1.0" @@ -1166,6 +991,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "bitcast.1.0" } id: 5 operand_ids: 6 @@ -1199,6 +1025,7 @@ gpu_executable { } } metadata { + scheduling_name: "out.1" } id: 3 operand_ids: 4 @@ -1262,26 +1089,17 @@ gpu_executable { value { instruction_ids: 0 instruction_ids: 8 - instruction_ids: 11 instruction_ids: 9 instruction_ids: 10 - instruction_ids: 2 instruction_ids: 1 instruction_ids: 7 - instruction_ids: 4 + instruction_ids: 2 instruction_ids: 6 + instruction_ids: 4 instruction_ids: 5 instruction_ids: 3 } } - sequences { - key: 4 - value { - instruction_ids: 17179869185 - instruction_ids: 17179869184 - instruction_ids: 17179869186 - } - } } input_output_alias { } @@ -1298,10 +1116,9 @@ gpu_executable { frontend_attributes { map { key: "fingerprint_before_lhs" - value: "5d6b0e05c28f8e5fa1780dcf27c57d60" + value: "7b2e4f79c5524b5cbc566da178aa7fda" } } - payloads: "{\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"native_emitter_backend_config\":{\"type\":\"NATIVE_EMITTER_TYPE_INVALID\",\"unroll_factor\":0},\"operation_queue_id\":\"0\",\"reification_cost\":[]}" payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" } config { @@ -1345,7 +1162,7 @@ gpu_executable { xla_cpu_enable_fast_math: false xla_gpu_enable_fast_min_max: false xla_force_host_platform_device_count: 1 - xla_gpu_enable_dynamic_slice_fusion: false + xla_gpu_enable_dynamic_slice_fusion: true xla_hlo_evaluator_use_fast_path: true xla_dump_hlo_as_html: false xla_cpu_fast_math_honor_nans: true @@ -1371,7 +1188,6 @@ gpu_executable { xla_gpu_use_runtime_fusion: false xla_dump_latency_hiding_schedule: false xla_dump_enable_mlir_pretty_form: true - xla_gpu_enable_latency_hiding_scheduler: false xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP xla_gpu_enable_triton_gemm: true xla_gpu_enable_cudnn_int8x32_convolution_reordering: true @@ -1384,14 +1200,10 @@ gpu_executable { xla_gpu_all_gather_combine_threshold_bytes: 31457287 xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 xla_gpu_enable_highest_priority_async_stream: true - xla_gpu_enable_pipelined_all_reduce: false xla_gpu_exhaustive_tiling_search: false xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 - xla_gpu_enable_pipelined_all_gather: false xla_gpu_redzone_padding_bytes: 8388608 - xla_gpu_enable_pipelined_reduce_scatter: true - xla_gpu_fused_attention_use_cudnn_rng: false xla_gpu_copy_insertion_use_region_analysis: false xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 xla_gpu_collect_cost_model_stats: false @@ -1435,6 +1247,7 @@ gpu_executable { xla_reduce_window_rewrite_base_length: 16 xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true xla_gpu_nccl_terminate_on_error: false xla_gpu_shard_autotuning: true xla_gpu_enable_approx_costly_collectives: false @@ -1443,7 +1256,7 @@ gpu_executable { xla_cmd_buffer_trace_cache_size: 16 xla_gpu_temp_buffer_use_separate_color: false xla_syntax_sugar_async_ops: false - xla_gpu_autotune_gemm_rtol: 0.1 + xla_gpu_autotune_gemm_rtol: 0.01 xla_enable_command_buffers_during_profiling: false xla_gpu_cudnn_gemm_max_plans: 5 xla_cpu_parallel_codegen_split_count: 32 @@ -1505,7 +1318,6 @@ gpu_executable { xla_cpu_generate_unique_c_style_kernel_entry_points: false xla_gpu_experimental_collective_cse_distance_threshold: 0 xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true - xla_cpu_use_fusion_emitters: true xla_gpu_experimental_collective_perf_table_path: "" xla_gpu_collective_permute_combine_threshold_bytes: 31457287 xla_unsupported_crash_on_hlo_pass_noop_change: false @@ -1514,7 +1326,6 @@ gpu_executable { xla_gpu_unsupported_enable_triton_multi_output_fusion: false xla_gpu_experimental_matmul_perf_table_path: "" xla_gpu_unsupported_enable_all_reduce_decomposer: false - xla_gpu_unsupported_use_all_reduce_one_shot_kernel: true xla_gpu_nccl_blocking_communicators: true xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 xla_gpu_first_collective_call_terminate_timeout_seconds: 40 @@ -1525,14 +1336,13 @@ gpu_executable { xla_detect_unstable_reductions: DETECTION_MODE_NONE xla_gpu_command_buffer_scheduling_mode: LHS xla_gpu_experimental_enable_nccl_symmetric_buffers: false - xla_gpu_experimental_autotuner_cache_dir: "" xla_disable_automatic_host_compute_offload: false xla_gpu_experimental_enable_fusion_autotuner: true xla_gpu_experimental_scaled_dot_with_triton: true xla_gpu_command_buffer_unroll_loops: false xla_cpu_experimental_onednn_custom_call: false - xla_gpu_experimental_use_raft_select_k: false xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true xla_cpu_collective_call_terminate_timeout_seconds: 40 xla_cpu_collective_call_warn_stuck_seconds: 20 xla_keep_shardings_after_spmd: false @@ -1542,21 +1352,32 @@ gpu_executable { xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE xla_cpu_enable_platform_dependent_math: true xla_gpu_detect_nan: DETECTION_MODE_NONE - xla_gpu_experimental_dump_gpu_executable: true xla_gpu_detect_inf: DETECTION_MODE_NONE xla_enable_enzyme_comms_opt: false xla_gpu_experimental_enable_buffer_saver_on_thunks: false xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE - xla_gpu_experimental_aot_compiled_thunks: true xla_enable_scoped_logging_timers: true xla_gpu_executable_embed_debug_info: true xla_cpu_collective_timeout_seconds: 1800 xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK xla_gpu_experimental_all_fusions_with_triton: false xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT xla_gpu_rocm_max_trace_events: 4194304 - xla_gpu_experimental_ragged_all_to_all_use_barrier: true xla_gpu_execution_terminate_timeout: "inf" xla_enable_hlo_sharding_v3: false xla_gpu_log_minmax: false @@ -1567,6 +1388,7 @@ gpu_executable { xla_gpu_experimental_max_unroll_factor: 32 xla_gpu_enable_pdl: true xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true xla_recognize_reduction_optimization_level: 0 xla_gpu_experimental_parallel_async_compute_limit: 2 xla_dump_buffer_assignment_analysis: true @@ -1586,12 +1408,12 @@ gpu_executable { xla_gpu_executable_num_compute_streams: 0 xla_gpu_executable_num_communication_streams: 0 xla_gpu_experimental_collective_start_as_early_as_possible: false - xla_cpu_experimental_enable_tiling_propagation: false + xla_cpu_experimental_enable_tiling_propagation: true xla_gpu_enable_dus_accumulator_zero_init_elimination: false xla_gpu_experimental_thunk_buffer_debug_module_outputs: false xla_cpu_ftz: true xla_gpu_enable_allocator_spatial_partitioning: true - xla_dump_compact_gte: false + xla_dump_compact_gte: true xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES xla_dump_hlo_as_riegeli: false xla_gpu_experimental_use_ragged_dot_grouped_gemm: true @@ -1600,6 +1422,35 @@ gpu_executable { xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED xla_gpu_enable_nccl_user_buffers_in_default_space: false xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" } static_device_assignment { replica_count: 2 @@ -1620,95 +1471,65 @@ gpu_executable { } } buffer_assignment { - logical_buffers { - id: 7 - size: 16 - defined_at { - instruction_id: 17179869186 - } - } - logical_buffers { - id: 8 - size: 16 - defined_at { - shape_index: 0 - instruction_id: 17179869186 - } - } logical_buffers { id: 9 - size: 16 - defined_at { - shape_index: 1 - instruction_id: 17179869186 - } - } - logical_buffers { - id: 12 size: 4 defined_at { } } logical_buffers { - id: 13 + id: 10 size: 8 defined_at { - instruction_id: 11 + instruction_id: 9 } } logical_buffers { - id: 14 + id: 11 size: 8 defined_at { instruction_id: 8 } } logical_buffers { - id: 15 + id: 12 size: 16 defined_at { - instruction_id: 9 + instruction_id: 10 } } logical_buffers { - id: 16 + id: 13 size: 16 defined_at { shape_index: 0 - instruction_id: 9 + instruction_id: 10 } } logical_buffers { - id: 17 + id: 14 size: 16 defined_at { shape_index: 1 - instruction_id: 9 - } - } - logical_buffers { - id: 18 - size: 16 - defined_at { instruction_id: 10 } } logical_buffers { - id: 19 + id: 15 size: 16 defined_at { instruction_id: 7 } } logical_buffers { - id: 20 + id: 16 size: 16 defined_at { instruction_id: 6 } } logical_buffers { - id: 21 + id: 17 size: 16 defined_at { instruction_id: 3 @@ -1718,12 +1539,12 @@ gpu_executable { size: 16 maybe_live_out: true assigned { - logical_buffer_id: 13 + logical_buffer_id: 10 size: 8 element_type: U32 } assigned { - logical_buffer_id: 19 + logical_buffer_id: 15 size: 16 element_type: U32 } @@ -1733,17 +1554,17 @@ gpu_executable { size: 16 maybe_live_out: true assigned { - logical_buffer_id: 8 - size: 16 + logical_buffer_id: 9 + size: 4 element_type: U32 } assigned { - logical_buffer_id: 12 - size: 4 + logical_buffer_id: 13 + size: 16 element_type: U32 } assigned { - logical_buffer_id: 20 + logical_buffer_id: 16 size: 16 element_type: F32 } @@ -1753,7 +1574,7 @@ gpu_executable { size: 16 maybe_live_out: true assigned { - logical_buffer_id: 21 + logical_buffer_id: 17 size: 16 element_type: TUPLE } @@ -1761,102 +1582,74 @@ gpu_executable { } buffer_allocations { index: 3 - size: 1544 + size: 520 assigned { - logical_buffer_id: 7 - offset: 1024 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 9 - offset: 1280 - size: 16 - element_type: F32 - } - assigned { - logical_buffer_id: 14 - offset: 1536 + logical_buffer_id: 11 + offset: 512 size: 8 element_type: F32 } assigned { - logical_buffer_id: 15 + logical_buffer_id: 12 size: 16 element_type: TUPLE } assigned { - logical_buffer_id: 16 + logical_buffer_id: 14 offset: 256 size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 17 - offset: 512 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 18 - offset: 768 - size: 16 - element_type: TUPLE + element_type: F32 } is_tuple: true } heap_simulator_traces { events { - buffer_id: 14 + buffer_id: 11 computation_name: "test_computation" instruction_name: "loop_broadcast_fusion" } events { - buffer_id: 9 - computation_name: "async_computation" - instruction_name: "allgather.2" + buffer_id: 14 + computation_name: "test_computation" + instruction_name: "allgather.3" } events { kind: FREE - buffer_id: 14 + buffer_id: 11 computation_name: "test_computation" instruction_name: "loop_broadcast_fusion" } events { kind: FREE - buffer_id: 9 - computation_name: "async_computation" - instruction_name: "allgather.2" + buffer_id: 14 + computation_name: "test_computation" + instruction_name: "allgather.3" } whole_module_simulation: true buffer_allocation_index: 3 } peak_buffers { - logical_buffer_ids: 19 + logical_buffer_ids: 15 } peak_buffers { index: 1 - logical_buffer_ids: 20 + logical_buffer_ids: 16 } peak_buffers { index: 2 - logical_buffer_ids: 21 + logical_buffer_ids: 17 } peak_buffers { index: 3 - logical_buffer_ids: 15 - logical_buffer_ids: 16 - logical_buffer_ids: 17 - logical_buffer_ids: 18 - logical_buffer_ids: 7 - logical_buffer_ids: 9 + logical_buffer_ids: 12 + logical_buffer_ids: 11 logical_buffer_ids: 14 } } gpu_compute_capability { cuda_compute_capability { - major: 6 - feature_extension: NONE + major: 9 + feature_extension: ACCELERATED_FEATURES } } module_name: "test" @@ -1887,92 +1680,21 @@ gpu_executable { values { size: 16 maybe_live_out: true - assigned { - logical_buffer_id: 13 - size: 8 - element_type: U32 - } - assigned { - logical_buffer_id: 19 - size: 16 - element_type: U32 - } } values { index: 1 size: 16 maybe_live_out: true - assigned { - logical_buffer_id: 8 - size: 16 - element_type: U32 - } - assigned { - logical_buffer_id: 12 - size: 4 - element_type: U32 - } - assigned { - logical_buffer_id: 20 - size: 16 - element_type: F32 - } } values { index: 2 size: 16 maybe_live_out: true - assigned { - logical_buffer_id: 21 - size: 16 - element_type: TUPLE - } is_tuple: true } values { index: 3 - size: 1544 - assigned { - logical_buffer_id: 7 - offset: 1024 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 9 - offset: 1280 - size: 16 - element_type: F32 - } - assigned { - logical_buffer_id: 14 - offset: 1536 - size: 8 - element_type: F32 - } - assigned { - logical_buffer_id: 15 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 16 - offset: 256 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 17 - offset: 512 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 18 - offset: 768 - size: 16 - element_type: TUPLE - } + size: 520 is_tuple: true } } @@ -2032,7 +1754,7 @@ gpu_executable { } args { slice { - offset: 1536 + offset: 512 size: 8 buffer_allocation_index: 3 element_type: F32 @@ -2056,7 +1778,7 @@ gpu_executable { name: "loop_broadcast_fusion" kernel_spec { cubin { - data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000@\010\000\000\000\000\000\000\300\005\000\000\000\000\000\000<\005<\000@\0008\000\003\000@\000\n\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.constant0.loop_broadcast_fusion\000.rel.nv.constant0.loop_broadcast_fusion\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.rel.nv.constant0.loop_broadcast_fusion\000.nv.constant0.loop_broadcast_fusion\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000loop_broadcast_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\266\000\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\332\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\000\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\300\000\000\000\000\000\000\000\004/\010\000\005\000\000\000\007\000\000\000\004\021\010\000\005\000\000\000\000\000\000\000\004\022\010\000\005\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\0010\000\000\001*\000\000\004\n\010\000\002\000\000\000@\001\020\000\003\031\020\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003\033\377\000\004\034\004\000x\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377s\000\000\000\000\000\000\000\000\000\000\021%\000\0056\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\007\000\376\000D\034\000\001\000\207\000\200\007\230L\002\000\007\005\200\007\230L\004\000\027\002\000\000\310\360\342\007\340\367\000\304?\010\003\000\027\005\200\007\230L\002\002\007\000\000\240\324\356\000\004\347\001\000\000(8\346\007\340\374\000t\034\004\004\004\'\005\000\201\030L\005\0007\005\000\010\020L\006\n\'\000\000\000\270\\\361\017\200\374\000\374\037\000\006\004\007\000\000 \334\356\000\017\007\000\000\000\260P\017\000\007\000\000\000\000\343\377\007\000\374\000\200\037\000\017\000\007\377\377\017@\342\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\340\007\000\374\000\200\037\000\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\005\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000E\001\000\000\000\000\000\000\033\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000`\002\000\000\000\000\000\000\220\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\002\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000N\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\003\000\000\000\000\000\000\\\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\332\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\003\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\366\000\000\000\013\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\216\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\003\000\000\000\000\000\000P\001\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\300\000\000\000\000\000\000\000\003\000\000\000\005\000\000\007 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\005\000\000\000@\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\240\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000@\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\240\013\000\000\000\000\000\000\240\010\000\000\000\000\000\000Z\rZ\000@\0008\000\005\000@\000\014\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_broadcast_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_broadcast_fusion\000.nv.constant0.loop_broadcast_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\216\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\275\000\000\000 \240\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\336\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000*\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000@\001\000\000\003\000\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004(\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\004\034\004\000\240\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\020\002\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\000\000\000\000D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\n\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\204\000\000\000\n\000\000\000\"\016\000\271z\004\000\000\202\000\000\000\n\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\206\000\000\000\n\000\000\000b\016\000\201y\002\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\342/\000Er\007\000\002\000\000\000\000\020 \000\000\312O\000\206y\000\004\007\000\000\000\004\031\020\014\000\342\017\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\024\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\211\001\000\000\000\000\000\000d\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\002\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\244\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\003\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\004\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000N\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000<\004\000\000\000\000\000\000`\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\324\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\234\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\302\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\004\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\216\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000 \002\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000" } arity: 2 kernel_name: "loop_broadcast_fusion" @@ -2089,6 +1811,7 @@ gpu_executable { } } } + use_pdl: true tma_metadata { } } @@ -2136,7 +1859,7 @@ gpu_executable { name: "wrapped_broadcast" kernel_spec { cubin { - data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\340\007\000\000\000\000\000\000`\005\000\000\000\000\000\000<\005<\000@\0008\000\003\000@\000\n\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.constant0.wrapped_broadcast\000.rel.nv.constant0.wrapped_broadcast\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.rel.nv.constant0.wrapped_broadcast\000.nv.constant0.wrapped_broadcast\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000wrapped_broadcast\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\246\000\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\306\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\342\000\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\361\000\000\000\022\020\t\000\000\000\000\000\000\000\000\000\200\000\000\000\000\000\000\000\004/\010\000\005\000\000\000\006\000\000\000\004\021\010\000\005\000\000\000\000\000\000\000\004\022\010\000\005\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\0010\000\000\001*\000\000\004\n\010\000\002\000\000\000@\001\020\000\003\031\020\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003\033\377\000\004\034\004\000p\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377s\000\000\000\000\000\000\000\000\000\000\021%\000\0056\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\007\000\376\000D\034\000\001\000\207\000\200\007\230L\002\000\007\005\200\007\230L\004\000\027\002\000\000\310\360\342\007\340\367\000\304?\010\003\000\027\005\200\007\230L\002\002\007\000\000\240\324\356\000\004\347\001\000\000(8\346\007\340\374\000\304\037\004\004\004\'\005\000\201\030L\005\0007\005\000\010\020L\002\004\007\000\000 \334\356\352\007\340\377\000\374\037\000\000\017\007\000\000\000\260P\017\000\007\000\000\000\000\343\017\000\207\377\377\017@\342\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\361\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0001\001\000\000\000\000\000\000\003\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\002\000\000\000\000\000\000\220\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\002\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000J\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\354\002\000\000\000\000\000\000\\\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\306\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\003\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\342\000\000\000\013\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000h\003\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\202\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\003\000\000\000\000\000\000P\001\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\004\000\000\000\000\000\000\200\000\000\000\000\000\000\000\003\000\000\000\005\000\000\006 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\005\000\000\000\340\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000x\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\350\001\000\000\000\000\000\000\350\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\340\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\240\013\000\000\000\000\000\000\240\010\000\000\000\000\000\000Z\rZ\000@\0008\000\005\000@\000\014\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.wrapped_broadcast\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000wrapped_broadcast\000.nv.constant0.wrapped_broadcast\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\261\000\000\000 \240\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\322\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\001\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\036\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\0000\001\000\000\003\000\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004$\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\004\034\004\000\220\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\020\002\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\000\000\000\000D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\n\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\204\000\000\000\n\000\000\000\"\016\000\271z\004\000\000\202\000\000\000\n\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\206\000\000\000\n\000\000\000b\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\312/\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\004\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000u\001\000\000\000\000\000\000P\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\002\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\230\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\003\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\003\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000J\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\004\000\000\000\000\000\000`\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000t\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\266\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\230\004\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\344\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000 \002\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000" } arity: 2 kernel_name: "wrapped_broadcast" @@ -2169,13 +1892,14 @@ gpu_executable { } } } + use_pdl: true tma_metadata { } } } thunks { thunk_info { - profile_annotation: "allgather.2" + profile_annotation: "allgather.3" thunk_id: 4 } all_gather_thunk { @@ -2227,7 +1951,7 @@ gpu_executable { element_count: 2 source_buffer { slice { - offset: 1536 + offset: 512 size: 8 buffer_allocation_index: 3 element_type: F32 @@ -2247,7 +1971,7 @@ gpu_executable { } destination_buffer { slice { - offset: 1280 + offset: 256 size: 16 buffer_allocation_index: 3 element_type: F32 @@ -2317,7 +2041,7 @@ gpu_executable { name: "loop_transpose_fusion_1" kernel_spec { cubin { - data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000`\010\000\000\000\000\000\000\340\005\000\000\000\000\000\000<\005<\000@\0008\000\003\000@\000\n\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.constant0.loop_transpose_fusion_1\000.rel.nv.constant0.loop_transpose_fusion_1\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.rel.nv.constant0.loop_transpose_fusion_1\000.nv.constant0.loop_transpose_fusion_1\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000loop_transpose_fusion_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\276\000\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\344\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\017\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\300\000\000\000\000\000\000\000\004/\010\000\005\000\000\000\006\000\000\000\004\021\010\000\005\000\000\000\000\000\000\000\004\022\010\000\005\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\0010\000\000\001*\000\000\004\n\010\000\002\000\000\000@\001\020\000\003\031\020\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003\033\377\000\004\034\004\000\220\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377s\000\000\000\000\000\000\000\000\000\000\021%\000\0056\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\007\340\343\000\230?\000\001\000\207\000\200\007\230L\005\000\027\002\000\000\310\360\000\005W\000\000\000\020\\\366\007\300\376\000\304\037\010\000\000\'\000\000\000\000\004\000\005\007\360\201\177\337[\003\000\347\001\000\000(8\346\007@\374\000\304\036\000\002\000\007\005\000\201\030L\003\003\027\005\000\010\020L\002\002\007\000\000\240\324\356\361\007\302\374\000\234\037\000\000\005\347\001\000\000(8\004\005\'\005\000\201\030L\005\0007\005\000\010\020L\361\007\341\377\000\374\037\000\002\004\007\000\000 \334\356\017\000\007\000\000\000\000\343\017\000\207\377\377\017@\342\340\007\000\374\000\200\037\000\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\017\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000O\001\000\000\000\000\000\000\'\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\002\000\000\000\000\000\000\220\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\003\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000P\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\003\000\000\000\000\000\000\\\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\344\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\003\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\001\000\000\013\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\003\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\224\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\270\003\000\000\000\000\000\000P\001\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \005\000\000\000\000\000\000\300\000\000\000\000\000\000\000\003\000\000\000\005\000\000\006 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\005\000\000\000`\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\270\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000(\002\000\000\000\000\000\000(\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000`\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\240\013\000\000\000\000\000\000\240\010\000\000\000\000\000\000Z\rZ\000@\0008\000\005\000@\000\014\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion_1\000.nv.constant0.loop_transpose_fusion_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\224\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\303\000\000\000 \240\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\344\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\001\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0000\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000H\001\000\000\003\000\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\020\002\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\000\000\000\000D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\n\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\204\000\000\000\n\000\000\000b\016\000\271z\004\000\000\202\000\000\000\n\000\000\000\342\017\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\206\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242\016\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\034\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\223\001\000\000\000\000\000\000n\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\252\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\003\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0000\004\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000P\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\004\000\000\000\000\000\000`\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\332\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\264\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\310\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\330\004\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\224\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000 \002\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000" } arity: 2 kernel_name: "loop_transpose_fusion_1" @@ -2350,6 +2074,7 @@ gpu_executable { } } } + use_pdl: true tma_metadata { } } @@ -2362,7 +2087,7 @@ gpu_executable { custom_kernel_thunk { args { slice { - offset: 1280 + offset: 256 size: 16 buffer_allocation_index: 3 element_type: F32 @@ -2405,7 +2130,7 @@ gpu_executable { name: "loop_transpose_fusion" kernel_spec { cubin { - data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000@\010\000\000\000\000\000\000\300\005\000\000\000\000\000\000<\005<\000@\0008\000\003\000@\000\n\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.constant0.loop_transpose_fusion\000.rel.nv.constant0.loop_transpose_fusion\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.rel.nv.constant0.loop_transpose_fusion\000.nv.constant0.loop_transpose_fusion\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000loop_transpose_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\266\000\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\332\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\000\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\300\000\000\000\000\000\000\000\004/\010\000\005\000\000\000\006\000\000\000\004\021\010\000\005\000\000\000\000\000\000\000\004\022\010\000\005\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\0010\000\000\001*\000\000\004\n\010\000\002\000\000\000@\001\020\000\003\031\020\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003\033\377\000\004\034\004\000\220\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377s\000\000\000\000\000\000\000\000\000\000\021%\000\0056\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\007\340\343\000\230?\000\001\000\207\000\200\007\230L\005\000\027\002\000\000\310\360\000\005W\000\000\000\020\\\366\007\300\376\000\304\037\010\000\000\'\000\000\000\000\004\000\005\007\360\201\177\337[\003\000\347\001\000\000(8\346\007@\374\000\304\036\000\002\000\007\005\000\201\030L\003\003\027\005\000\010\020L\002\002\007\000\000\240\324\356\361\007\302\374\000\234\037\000\000\005\347\001\000\000(8\004\005\'\005\000\201\030L\005\0007\005\000\010\020L\361\007\341\377\000\374\037\000\002\004\007\000\000 \334\356\017\000\007\000\000\000\000\343\017\000\207\377\377\017@\342\340\007\000\374\000\200\037\000\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\005\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000E\001\000\000\000\000\000\000\033\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000`\002\000\000\000\000\000\000\220\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\002\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000N\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\003\000\000\000\000\000\000\\\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\332\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\003\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\366\000\000\000\013\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\216\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\003\000\000\000\000\000\000P\001\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\300\000\000\000\000\000\000\000\003\000\000\000\005\000\000\006 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\005\000\000\000@\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\240\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000@\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\240\013\000\000\000\000\000\000\240\010\000\000\000\000\000\000Z\rZ\000@\0008\000\005\000@\000\014\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion\000.nv.constant0.loop_transpose_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\216\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\275\000\000\000 \240\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\336\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000*\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000@\001\000\000\003\000\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\020\002\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\000\000\000\000D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\n\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\204\000\000\000\n\000\000\000b\016\000\271z\004\000\000\202\000\000\000\n\000\000\000\342\017\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\206\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242\016\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\024\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\211\001\000\000\000\000\000\000d\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\002\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\244\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\003\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\004\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000N\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000<\004\000\000\000\000\000\000`\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\324\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\234\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\302\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\004\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\216\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000 \002\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000" } arity: 2 kernel_name: "loop_transpose_fusion" @@ -2438,6 +2163,7 @@ gpu_executable { } } } + use_pdl: true tma_metadata { } } @@ -2446,9 +2172,9 @@ gpu_executable { platform_name: "CUDA" cuda_platform_version { cuda_toolkit_version: "12.8.0" - cudnn_version: "9.10.0" + cudnn_version: "9.10.2" cub_version: "3.1.2" } } - buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 0B\n maybe_live_out allocation: 48B\n preallocated temp allocation: 1.5KiB\n preallocated temp fragmentation: 240B (15.54%)\n total allocation: 1.6KiB\n total fragmentation: 1.4KiB (91.96%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: (u32[2,2], f32[2,2])\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 16B\n\t\tXLA Label: async-done\n\t\tShape: (u32[2,2], f32[2,2])\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: async-start\n\t\tShape: (u32[2,2], f32[2,2])\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 16B\n\t\tXLA Label: async-start\n\t\tShape: (u32[2,1], f32[2,1])\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 16B\n\t\tXLA Label: async-start\n\t\tShape: ((u32[2,1], f32[2,1]), (u32[2,2], f32[2,2]))\n\t\t==========================\n\n\tBuffer 7:\n\t\tSize: 16B\n\t\tXLA Label: tuple\n\t\tShape: (u32[4], f32[4])\n\t\t==========================\n\n\tBuffer 8:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 9:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: u32[2,2]\n\t\t==========================\n\n\tBuffer 10:\n\t\tSize: 8B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,1]\n\t\t==========================\n" + buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 0B\n maybe_live_out allocation: 48B\n preallocated temp allocation: 520B\n preallocated temp fragmentation: 240B (46.15%)\n total allocation: 568B\n total fragmentation: 504B (88.73%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: (u32[2,2], f32[2,2])\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 16B\n\t\tXLA Label: tuple\n\t\tShape: (u32[4], f32[4])\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: u32[2,2]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 8B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,1]\n\t\t==========================\n" } diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu/v2/ReduceScatter.pbtxt b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_h100/v1/ReduceScatter.pbtxt similarity index 67% rename from third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu/v2/ReduceScatter.pbtxt rename to third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_h100/v1/ReduceScatter.pbtxt index 24fe07b3ad1846..89dadf3e45ea87 100644 --- a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu/v2/ReduceScatter.pbtxt +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_h100/v1/ReduceScatter.pbtxt @@ -1,9 +1,9 @@ -# proto-file: third_party/tensorflow/compiler/xla/tools/split_proto/split_proto_cli.proto -# proto-message: xla.split_proto_cli.DeserializedSplitExecutableAndOptions - +# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto +# proto-message: xla.HumanReadableAotExecutable executable_and_options { compile_options { executable_build_options { + device_ordinal: -1 result_layout { element_type: U32 dimensions: 4 @@ -27,7 +27,7 @@ executable_and_options { xla_cpu_enable_fast_math: false xla_gpu_enable_fast_min_max: false xla_force_host_platform_device_count: 1 - xla_gpu_enable_dynamic_slice_fusion: false + xla_gpu_enable_dynamic_slice_fusion: true xla_hlo_evaluator_use_fast_path: true xla_dump_hlo_as_html: false xla_cpu_fast_math_honor_nans: true @@ -53,7 +53,6 @@ executable_and_options { xla_gpu_use_runtime_fusion: false xla_dump_latency_hiding_schedule: false xla_dump_enable_mlir_pretty_form: true - xla_gpu_enable_latency_hiding_scheduler: false xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP xla_gpu_enable_triton_gemm: true xla_gpu_enable_cudnn_int8x32_convolution_reordering: true @@ -66,14 +65,10 @@ executable_and_options { xla_gpu_all_gather_combine_threshold_bytes: 31457287 xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 xla_gpu_enable_highest_priority_async_stream: true - xla_gpu_enable_pipelined_all_reduce: false xla_gpu_exhaustive_tiling_search: false xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 - xla_gpu_enable_pipelined_all_gather: false xla_gpu_redzone_padding_bytes: 8388608 - xla_gpu_enable_pipelined_reduce_scatter: true - xla_gpu_fused_attention_use_cudnn_rng: false xla_gpu_copy_insertion_use_region_analysis: false xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 xla_gpu_collect_cost_model_stats: false @@ -117,6 +112,7 @@ executable_and_options { xla_reduce_window_rewrite_base_length: 16 xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true xla_gpu_nccl_terminate_on_error: false xla_gpu_shard_autotuning: true xla_gpu_enable_approx_costly_collectives: false @@ -125,7 +121,7 @@ executable_and_options { xla_cmd_buffer_trace_cache_size: 16 xla_gpu_temp_buffer_use_separate_color: false xla_syntax_sugar_async_ops: false - xla_gpu_autotune_gemm_rtol: 0.1 + xla_gpu_autotune_gemm_rtol: 0.01 xla_enable_command_buffers_during_profiling: false xla_gpu_cudnn_gemm_max_plans: 5 xla_cpu_parallel_codegen_split_count: 32 @@ -187,7 +183,6 @@ executable_and_options { xla_cpu_generate_unique_c_style_kernel_entry_points: false xla_gpu_experimental_collective_cse_distance_threshold: 0 xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true - xla_cpu_use_fusion_emitters: true xla_gpu_experimental_collective_perf_table_path: "" xla_gpu_collective_permute_combine_threshold_bytes: 31457287 xla_unsupported_crash_on_hlo_pass_noop_change: false @@ -196,7 +191,6 @@ executable_and_options { xla_gpu_unsupported_enable_triton_multi_output_fusion: false xla_gpu_experimental_matmul_perf_table_path: "" xla_gpu_unsupported_enable_all_reduce_decomposer: false - xla_gpu_unsupported_use_all_reduce_one_shot_kernel: true xla_gpu_nccl_blocking_communicators: true xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 xla_gpu_first_collective_call_terminate_timeout_seconds: 40 @@ -207,14 +201,13 @@ executable_and_options { xla_detect_unstable_reductions: DETECTION_MODE_NONE xla_gpu_command_buffer_scheduling_mode: LHS xla_gpu_experimental_enable_nccl_symmetric_buffers: false - xla_gpu_experimental_autotuner_cache_dir: "" xla_disable_automatic_host_compute_offload: false xla_gpu_experimental_enable_fusion_autotuner: true xla_gpu_experimental_scaled_dot_with_triton: true xla_gpu_command_buffer_unroll_loops: false xla_cpu_experimental_onednn_custom_call: false - xla_gpu_experimental_use_raft_select_k: false xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true xla_cpu_collective_call_terminate_timeout_seconds: 40 xla_cpu_collective_call_warn_stuck_seconds: 20 xla_keep_shardings_after_spmd: false @@ -224,21 +217,32 @@ executable_and_options { xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE xla_cpu_enable_platform_dependent_math: true xla_gpu_detect_nan: DETECTION_MODE_NONE - xla_gpu_experimental_dump_gpu_executable: true xla_gpu_detect_inf: DETECTION_MODE_NONE xla_enable_enzyme_comms_opt: false xla_gpu_experimental_enable_buffer_saver_on_thunks: false xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE - xla_gpu_experimental_aot_compiled_thunks: true xla_enable_scoped_logging_timers: true xla_gpu_executable_embed_debug_info: true xla_cpu_collective_timeout_seconds: 1800 xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK xla_gpu_experimental_all_fusions_with_triton: false xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT xla_gpu_rocm_max_trace_events: 4194304 - xla_gpu_experimental_ragged_all_to_all_use_barrier: true xla_gpu_execution_terminate_timeout: "inf" xla_enable_hlo_sharding_v3: false xla_gpu_log_minmax: false @@ -249,6 +253,7 @@ executable_and_options { xla_gpu_experimental_max_unroll_factor: 32 xla_gpu_enable_pdl: true xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true xla_recognize_reduction_optimization_level: 0 xla_gpu_experimental_parallel_async_compute_limit: 2 xla_dump_buffer_assignment_analysis: true @@ -268,12 +273,12 @@ executable_and_options { xla_gpu_executable_num_compute_streams: 0 xla_gpu_executable_num_communication_streams: 0 xla_gpu_experimental_collective_start_as_early_as_possible: false - xla_cpu_experimental_enable_tiling_propagation: false + xla_cpu_experimental_enable_tiling_propagation: true xla_gpu_enable_dus_accumulator_zero_init_elimination: false xla_gpu_experimental_thunk_buffer_debug_module_outputs: false xla_cpu_ftz: true xla_gpu_enable_allocator_spatial_partitioning: true - xla_dump_compact_gte: false + xla_dump_compact_gte: true xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES xla_dump_hlo_as_riegeli: false xla_gpu_experimental_use_ragged_dot_grouped_gemm: true @@ -282,6 +287,35 @@ executable_and_options { xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED xla_gpu_enable_nccl_user_buffers_in_default_space: false xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" } num_replicas: 2 num_partitions: 1 @@ -299,9 +333,9 @@ executable_and_options { allow_spmd_sharding_propagation_to_parameters: false process_count: 1 memory_fitting_level: EFFORT_O2 - slice_size: 2 } } + pjrt_client_name: "PjRtStreamExecutorClient" } gpu_executable { hlo_module_with_config { @@ -320,6 +354,7 @@ gpu_executable { } } metadata { + scheduling_name: "lhs" } frontend_attributes { } @@ -336,6 +371,7 @@ gpu_executable { } } metadata { + scheduling_name: "rhs" } parameter_number: 1 id: 1 @@ -354,6 +390,7 @@ gpu_executable { } } metadata { + scheduling_name: "add.1" } id: 2 operand_ids: 0 @@ -399,9 +436,10 @@ gpu_executable { } } metadata { + scheduling_name: "param_2.2" } parameter_number: 2 - id: 12884901893 + id: 8589934597 frontend_attributes { } statistics_viz { @@ -427,14 +465,14 @@ gpu_executable { } u32s: 0 } - id: 12884901894 + id: 8589934598 frontend_attributes { } statistics_viz { } } instructions { - name: "p.4" + name: "p.2" opcode: "compare" shape { element_type: PRED @@ -443,19 +481,21 @@ gpu_executable { } } metadata { + scheduling_name: "p.2" } - id: 12884901892 - operand_ids: 12884901893 - operand_ids: 12884901894 + id: 8589934596 + operand_ids: 8589934597 + operand_ids: 8589934598 comparison_direction: "EQ" frontend_attributes { } comparison_type: "UNSIGNED" statistics_viz { } + comparison_order: "TOTAL" } instructions { - name: "pb.4" + name: "pb.2" opcode: "broadcast" shape { element_type: PRED @@ -467,16 +507,17 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "pb.2" } - id: 12884901889 - operand_ids: 12884901892 + id: 8589934593 + operand_ids: 8589934596 frontend_attributes { } statistics_viz { } } instructions { - name: "param_0.4" + name: "param_0.1" opcode: "parameter" shape { element_type: U32 @@ -488,15 +529,16 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "param_0.1" } - id: 12884901890 + id: 8589934594 frontend_attributes { } statistics_viz { } } instructions { - name: "param_1.3" + name: "param_1.1" opcode: "parameter" shape { element_type: U32 @@ -508,9 +550,10 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "param_1.1" } parameter_number: 1 - id: 12884901891 + id: 8589934595 frontend_attributes { } statistics_viz { @@ -529,11 +572,12 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "data.2" } - id: 12884901888 - operand_ids: 12884901889 - operand_ids: 12884901890 - operand_ids: 12884901891 + id: 8589934592 + operand_ids: 8589934593 + operand_ids: 8589934594 + operand_ids: 8589934595 frontend_attributes { } statistics_viz { @@ -573,85 +617,14 @@ gpu_executable { } is_dynamic_dimension: false } - parameter_names: "param_0.4" - parameter_names: "param_1.3" + parameter_names: "param_0.1" + parameter_names: "param_1.1" parameter_names: "param_2.2" } - id: 3 - root_id: 12884901888 + id: 2 + root_id: 8589934592 is_fusion_computation: true } - computations { - name: "async_computation" - instructions { - name: "param_0.5" - opcode: "parameter" - shape { - element_type: U32 - dimensions: 8 - layout { - minor_to_major: 0 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - } - metadata { - } - id: 21474836480 - frontend_attributes { - } - statistics_viz { - } - } - instructions { - name: "ars.2" - opcode: "reduce-scatter" - shape { - element_type: U32 - dimensions: 4 - layout { - minor_to_major: 0 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - } - metadata { - } - dimensions: 0 - id: 21474836481 - operand_ids: 21474836480 - called_computation_ids: 0 - frontend_attributes { - } - statistics_viz { - } - collective_device_list { - } - } - program_shape { - parameters { - element_type: U32 - dimensions: 8 - layout { - minor_to_major: 0 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - } - result { - element_type: U32 - dimensions: 4 - layout { - minor_to_major: 0 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - } - parameter_names: "param_0.5" - } - id: 5 - root_id: 21474836481 - } computations { name: "main" instructions { @@ -742,6 +715,7 @@ gpu_executable { } } metadata { + scheduling_name: "id.1" } id: 4294967296 frontend_attributes { @@ -762,64 +736,22 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "loop_select_fusion" } fusion_kind: "kLoop" id: 4294967299 operand_ids: 4294967297 operand_ids: 4294967298 operand_ids: 4294967296 - called_computation_ids: 3 + called_computation_ids: 2 frontend_attributes { } statistics_viz { } - backend_config_payload { - id: 0 - } } instructions { - name: "reduce-scatter-start" - opcode: "async-start" - shape { - element_type: TUPLE - tuple_shapes { - element_type: TUPLE - tuple_shapes { - element_type: U32 - dimensions: 8 - layout { - minor_to_major: 0 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - } - } - tuple_shapes { - element_type: U32 - dimensions: 4 - layout { - minor_to_major: 0 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - } - } - metadata { - } - id: 4294967300 - operand_ids: 4294967299 - called_computation_ids: 5 - frontend_attributes { - } - statistics_viz { - } - backend_config_payload { - id: 1 - } - } - instructions { - name: "reduce-scatter-done" - opcode: "async-done" + name: "ars.3" + opcode: "reduce-scatter" shape { element_type: U32 dimensions: 4 @@ -830,13 +762,21 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "reduce-scatter-start" } - id: 4294967301 - operand_ids: 4294967300 + dimensions: 0 + id: 4294967300 + operand_ids: 4294967299 + called_computation_ids: 0 frontend_attributes { } statistics_viz { } + collective_device_list { + } + backend_config_payload { + id: 0 + } } program_shape { result { @@ -850,7 +790,7 @@ gpu_executable { } } id: 1 - root_id: 4294967301 + root_id: 4294967300 } host_program_shape { result { @@ -881,14 +821,6 @@ gpu_executable { instruction_ids: 4294967297 instruction_ids: 4294967299 instruction_ids: 4294967300 - instruction_ids: 4294967301 - } - } - sequences { - key: 5 - value { - instruction_ids: 21474836480 - instruction_ids: 21474836481 } } } @@ -907,10 +839,9 @@ gpu_executable { frontend_attributes { map { key: "fingerprint_before_lhs" - value: "f1e11fca63da83b4e8c82bd1944aded5" + value: "c36d8870be58d238935bebefbbc5b827" } } - payloads: "{\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"native_emitter_backend_config\":{\"type\":\"NATIVE_EMITTER_TYPE_INVALID\",\"unroll_factor\":0},\"operation_queue_id\":\"0\",\"reification_cost\":[]}" payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" } config { @@ -942,7 +873,7 @@ gpu_executable { xla_cpu_enable_fast_math: false xla_gpu_enable_fast_min_max: false xla_force_host_platform_device_count: 1 - xla_gpu_enable_dynamic_slice_fusion: false + xla_gpu_enable_dynamic_slice_fusion: true xla_hlo_evaluator_use_fast_path: true xla_dump_hlo_as_html: false xla_cpu_fast_math_honor_nans: true @@ -968,7 +899,6 @@ gpu_executable { xla_gpu_use_runtime_fusion: false xla_dump_latency_hiding_schedule: false xla_dump_enable_mlir_pretty_form: true - xla_gpu_enable_latency_hiding_scheduler: false xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP xla_gpu_enable_triton_gemm: true xla_gpu_enable_cudnn_int8x32_convolution_reordering: true @@ -981,14 +911,10 @@ gpu_executable { xla_gpu_all_gather_combine_threshold_bytes: 31457287 xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 xla_gpu_enable_highest_priority_async_stream: true - xla_gpu_enable_pipelined_all_reduce: false xla_gpu_exhaustive_tiling_search: false xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 - xla_gpu_enable_pipelined_all_gather: false xla_gpu_redzone_padding_bytes: 8388608 - xla_gpu_enable_pipelined_reduce_scatter: true - xla_gpu_fused_attention_use_cudnn_rng: false xla_gpu_copy_insertion_use_region_analysis: false xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 xla_gpu_collect_cost_model_stats: false @@ -1032,6 +958,7 @@ gpu_executable { xla_reduce_window_rewrite_base_length: 16 xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true xla_gpu_nccl_terminate_on_error: false xla_gpu_shard_autotuning: true xla_gpu_enable_approx_costly_collectives: false @@ -1040,7 +967,7 @@ gpu_executable { xla_cmd_buffer_trace_cache_size: 16 xla_gpu_temp_buffer_use_separate_color: false xla_syntax_sugar_async_ops: false - xla_gpu_autotune_gemm_rtol: 0.1 + xla_gpu_autotune_gemm_rtol: 0.01 xla_enable_command_buffers_during_profiling: false xla_gpu_cudnn_gemm_max_plans: 5 xla_cpu_parallel_codegen_split_count: 32 @@ -1102,7 +1029,6 @@ gpu_executable { xla_cpu_generate_unique_c_style_kernel_entry_points: false xla_gpu_experimental_collective_cse_distance_threshold: 0 xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true - xla_cpu_use_fusion_emitters: true xla_gpu_experimental_collective_perf_table_path: "" xla_gpu_collective_permute_combine_threshold_bytes: 31457287 xla_unsupported_crash_on_hlo_pass_noop_change: false @@ -1111,7 +1037,6 @@ gpu_executable { xla_gpu_unsupported_enable_triton_multi_output_fusion: false xla_gpu_experimental_matmul_perf_table_path: "" xla_gpu_unsupported_enable_all_reduce_decomposer: false - xla_gpu_unsupported_use_all_reduce_one_shot_kernel: true xla_gpu_nccl_blocking_communicators: true xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 xla_gpu_first_collective_call_terminate_timeout_seconds: 40 @@ -1122,14 +1047,13 @@ gpu_executable { xla_detect_unstable_reductions: DETECTION_MODE_NONE xla_gpu_command_buffer_scheduling_mode: LHS xla_gpu_experimental_enable_nccl_symmetric_buffers: false - xla_gpu_experimental_autotuner_cache_dir: "" xla_disable_automatic_host_compute_offload: false xla_gpu_experimental_enable_fusion_autotuner: true xla_gpu_experimental_scaled_dot_with_triton: true xla_gpu_command_buffer_unroll_loops: false xla_cpu_experimental_onednn_custom_call: false - xla_gpu_experimental_use_raft_select_k: false xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true xla_cpu_collective_call_terminate_timeout_seconds: 40 xla_cpu_collective_call_warn_stuck_seconds: 20 xla_keep_shardings_after_spmd: false @@ -1139,21 +1063,32 @@ gpu_executable { xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE xla_cpu_enable_platform_dependent_math: true xla_gpu_detect_nan: DETECTION_MODE_NONE - xla_gpu_experimental_dump_gpu_executable: true xla_gpu_detect_inf: DETECTION_MODE_NONE xla_enable_enzyme_comms_opt: false xla_gpu_experimental_enable_buffer_saver_on_thunks: false xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE - xla_gpu_experimental_aot_compiled_thunks: true xla_enable_scoped_logging_timers: true xla_gpu_executable_embed_debug_info: true xla_cpu_collective_timeout_seconds: 1800 xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK xla_gpu_experimental_all_fusions_with_triton: false xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT xla_gpu_rocm_max_trace_events: 4194304 - xla_gpu_experimental_ragged_all_to_all_use_barrier: true xla_gpu_execution_terminate_timeout: "inf" xla_enable_hlo_sharding_v3: false xla_gpu_log_minmax: false @@ -1164,6 +1099,7 @@ gpu_executable { xla_gpu_experimental_max_unroll_factor: 32 xla_gpu_enable_pdl: true xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true xla_recognize_reduction_optimization_level: 0 xla_gpu_experimental_parallel_async_compute_limit: 2 xla_dump_buffer_assignment_analysis: true @@ -1183,12 +1119,12 @@ gpu_executable { xla_gpu_executable_num_compute_streams: 0 xla_gpu_executable_num_communication_streams: 0 xla_gpu_experimental_collective_start_as_early_as_possible: false - xla_cpu_experimental_enable_tiling_propagation: false + xla_cpu_experimental_enable_tiling_propagation: true xla_gpu_enable_dus_accumulator_zero_init_elimination: false xla_gpu_experimental_thunk_buffer_debug_module_outputs: false xla_cpu_ftz: true xla_gpu_enable_allocator_spatial_partitioning: true - xla_dump_compact_gte: false + xla_dump_compact_gte: true xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES xla_dump_hlo_as_riegeli: false xla_gpu_experimental_use_ragged_dot_grouped_gemm: true @@ -1197,6 +1133,35 @@ gpu_executable { xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED xla_gpu_enable_nccl_user_buffers_in_default_space: false xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" } static_device_assignment { replica_count: 2 @@ -1238,58 +1203,43 @@ gpu_executable { } logical_buffers { id: 10 - size: 16 - defined_at { - instruction_id: 21474836481 - } - } - logical_buffers { - id: 11 size: 32 defined_at { instruction_id: 4294967297 } } logical_buffers { - id: 12 + id: 11 size: 32 defined_at { instruction_id: 4294967298 } } logical_buffers { - id: 13 + id: 12 size: 4 defined_at { instruction_id: 4294967296 } } logical_buffers { - id: 14 + id: 13 size: 32 defined_at { instruction_id: 4294967299 } } logical_buffers { - id: 15 + id: 14 size: 16 defined_at { instruction_id: 4294967300 } } - logical_buffers { - id: 16 - size: 8 - defined_at { - shape_index: 0 - instruction_id: 4294967300 - } - } buffer_allocations { size: 32 assigned { - logical_buffer_id: 11 + logical_buffer_id: 10 size: 32 element_type: U32 } @@ -1299,7 +1249,7 @@ gpu_executable { index: 1 size: 32 assigned { - logical_buffer_id: 12 + logical_buffer_id: 11 size: 32 element_type: U32 } @@ -1310,13 +1260,13 @@ gpu_executable { size: 16 maybe_live_out: true assigned { - logical_buffer_id: 10 - size: 16 + logical_buffer_id: 12 + size: 4 element_type: U32 } assigned { - logical_buffer_id: 13 - size: 4 + logical_buffer_id: 14 + size: 16 element_type: U32 } } @@ -1335,7 +1285,6 @@ gpu_executable { size: 4 is_thread_local: true assigned { - logical_buffer_id: 1 size: 4 element_type: U32 } @@ -1345,56 +1294,45 @@ gpu_executable { size: 4 is_thread_local: true assigned { + logical_buffer_id: 1 size: 4 element_type: U32 } } buffer_allocations { index: 6 - size: 528 + size: 32 assigned { - logical_buffer_id: 14 + logical_buffer_id: 13 size: 32 element_type: U32 } - assigned { - logical_buffer_id: 15 - offset: 512 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 16 - offset: 256 - size: 8 - element_type: TUPLE - } } heap_simulator_traces { events { - buffer_id: 14 + buffer_id: 13 computation_name: "main" instruction_name: "loop_select_fusion" } events { kind: FREE - buffer_id: 14 + buffer_id: 13 computation_name: "main" instruction_name: "loop_select_fusion" } whole_module_simulation: true - buffer_allocation_index: 5 + buffer_allocation_index: 3 } peak_buffers { - logical_buffer_ids: 11 + logical_buffer_ids: 10 } peak_buffers { index: 1 - logical_buffer_ids: 12 + logical_buffer_ids: 11 } peak_buffers { index: 2 - logical_buffer_ids: 10 + logical_buffer_ids: 14 } peak_buffers { index: 3 @@ -1402,24 +1340,22 @@ gpu_executable { } peak_buffers { index: 4 - logical_buffer_ids: 1 + logical_buffer_ids: 0 } peak_buffers { index: 5 - logical_buffer_ids: 0 + logical_buffer_ids: 1 } peak_buffers { index: 6 - logical_buffer_ids: 14 - logical_buffer_ids: 16 - logical_buffer_ids: 15 + logical_buffer_ids: 13 } } - binary: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000@\005\000\000\000\000\000\000\000\003\000\000\000\000\000\000<\005<\000@\0008\000\004\000@\000\t\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.nv.global\000.nv.constant4\000.rel.nv.constant4\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.nv.global\000.nv.constant4\000.rel.nv.constant4\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000buffer_for_c1_1\000buffer_for_c0_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000=\000\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000y\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\000\000\000\021\000\010\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\230\000\000\000\021\000\010\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377s\000\000\000\000\000\000\000\000\000\000\021%\000\0056\010\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\210\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000]\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\002\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000y\000\000\000\013\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\002\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000K\000\000\000\t\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\002\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\007\000\000\000\010\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000=\000\000\000\001\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000h\002\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\003\000\000\000\000\000\000@\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\005\000\000\000@\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000h\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000@\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + binary: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000\000\004\000\000\000\000\000\000Z\rZ\000@\0008\000\004\000@\000\n\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.nv.shared.reserved.0\000.nv.global\000.nv.constant4\000\000\000\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant4\000\000\000\000.nv.callgraph\000.nv.prototype\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.global\000.nv.constant4\000\000\000\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant.pic\000.nv.callgraph\000.nv.prototype\000buffer_for_c1_1\000buffer_for_c0_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000a\000\000\000 \240\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\215\000\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\326\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\362\000\000\000\021\000\t\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\002\001\000\000\021\000\t\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\010\000\000\000\000\000\000\000\002\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\270\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\001\000\000\000\000\000\000\022\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\002\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000d\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\330\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\234\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\330\002\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\206\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\370\002\000\000\000\000\000\0000\000\000\000\000\000\000\000\003\000\000\000\007\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000S\000\000\000\001\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000(\003\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000@\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000(\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\0008\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\000\000\000\000\000\000\000\010\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000" gpu_compute_capability { cuda_compute_capability { - major: 6 - feature_extension: NONE + major: 9 + feature_extension: ACCELERATED_FEATURES } } module_name: "test" @@ -1437,87 +1373,36 @@ gpu_executable { buffer_allocations { values { size: 32 - assigned { - logical_buffer_id: 11 - size: 32 - element_type: U32 - } is_constant: true } values { index: 1 size: 32 - assigned { - logical_buffer_id: 12 - size: 32 - element_type: U32 - } is_constant: true } values { index: 2 size: 16 maybe_live_out: true - assigned { - logical_buffer_id: 10 - size: 16 - element_type: U32 - } - assigned { - logical_buffer_id: 13 - size: 4 - element_type: U32 - } } values { index: 3 size: 4 is_thread_local: true - assigned { - logical_buffer_id: 2 - size: 4 - element_type: U32 - } } values { index: 4 size: 4 is_thread_local: true - assigned { - logical_buffer_id: 1 - size: 4 - element_type: U32 - } } values { index: 5 size: 4 is_thread_local: true - assigned { - size: 4 - element_type: U32 - } } values { index: 6 - size: 528 - assigned { - logical_buffer_id: 14 - size: 32 - element_type: U32 - } - assigned { - logical_buffer_id: 15 - offset: 512 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 16 - offset: 256 - size: 8 - element_type: TUPLE - } + size: 32 } } output_info_map { @@ -1621,7 +1506,7 @@ gpu_executable { name: "loop_select_fusion" kernel_spec { cubin { - data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\000\006\000\000\000\000\000\000<\005<\000@\0008\000\003\000@\000\n\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.constant0.loop_select_fusion\000.rel.nv.constant0.loop_select_fusion\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.rel.nv.constant0.loop_select_fusion\000.nv.constant0.loop_select_fusion\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000loop_select_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\252\000\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\313\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\347\000\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\000\000\000\022\020\t\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\004/\010\000\005\000\000\000\007\000\000\000\004\021\010\000\005\000\000\000\000\000\000\000\004\022\010\000\005\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\0010\000\000\001*\000\000\004\n\010\000\002\000\000\000@\001 \000\003\031 \000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003\033\377\000\004\034\004\000\270\000\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377s\000\000\000\000\000\000\000\000\000\000\021%\000\0056\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\007\000\376\000D\034\000\001\000\207\000\200\007\230L\002\000G\005\200\007\230L\000\000\027\002\000\000\310\360\342\007 \366\000\204\037\000\003\000W\005\200\007\230L\002\002\007\000\000\240\324\356\004\000\007\005\200\007\230L\377\007 \376A\304\037\000\005\000\027\005\200\007\230L\006\000\'\000\000\000H8\000\000\347\001\000\000(8\361\007\251\374\000\230\037\000\004\004\'\005\000\001DK\005\0057\005\000\001DK\004\004g\000\000\200\020\\\342\007 \366\000\230\037\000\005\005\007\000\000\010\020\\\004\004\007\000\000\240\324\356\002\006g\005\000\200\020L\347\007 \376 \374\037\000\003\000w\005\000\010\020L\004\002\007\000\000 \334\356\017\000\007\000\000\000\000\343\377\007\000\374\000\200\037\000\017\000\007\377\377\017@\342\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\340\007\000\374\000\200\037\000\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\366\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0006\001\000\000\000\000\000\000\t\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\002\000\000\000\000\000\000\220\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\002\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000K\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\364\002\000\000\000\000\000\000|\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\313\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\003\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\347\000\000\000\013\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\205\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\003\000\000\000\000\000\000`\001\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\000\001\000\000\000\000\000\000\003\000\000\000\005\000\000\007 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\005\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\240\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000`\002\000\000\000\000\000\000`\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\0000\014\000\000\000\000\000\0000\t\000\000\000\000\000\000Z\rZ\000@\0008\000\005\000@\000\014\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_select_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_select_fusion\000.nv.constant0.loop_select_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\205\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\264\000\000\000 \240\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\325\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\001\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000!\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\0004\001\000\000\003\000\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\004<\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\004/\010\000\006\000\000\000\014\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\004\034\004\000\360\000\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000\003\031 \000\004\n\010\000\007\000\000\000\020\002 \000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\000\000\000\000D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\n\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\210\000\000\000\n\000\000\000\"\016\000\271z\004\000\000\202\000\000\000\n\000\000\000\316\017\000\202{\004\377\000\204\000\000\000\n\000\000\000b\016\000\031y\t\000\000\000\000\000\000!\000\000\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\314\017\000\202{\006\377\000\212\000\000\000\n\000\000\000\342\016\000\201y\002\002\004\000\000\000\000\231\036\014\000$\037\000\014r\000\002\377\000\000\000pP\360\003\000\332\017\001\202\013\004\377\000\206\000\000\000\010\000\000\000\260.\000\202\013\005\377\000\207\000\000\000\010\000\000\000\244\016\000%x\004\t\004\000\000\000\004\000\216\007\000\314O\000\201y\005\004\004\000\000\000\000\231\036\014\000\242\016\000%x\006\t\004\000\000\000\006\000\216\007\000\312\217\000\206y\000\006\005\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000z\001\000\000\000\000\000\000U\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\002\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\233\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\370\003\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000K\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\034\004\000\000\000\000\000\000\200\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\313\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\234\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\271\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\004\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\000\002\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\205\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\347\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\0000\002\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\0000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\0000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\000\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0000\002\000\000\000\000\000\0000\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000" } arity: 4 kernel_name: "loop_select_fusion" @@ -1666,13 +1551,14 @@ gpu_executable { } } } + use_pdl: true tma_metadata { } } } thunks { thunk_info { - profile_annotation: "ars.2" + profile_annotation: "ars.3" thunk_id: 3 } reduce_scatter_thunk { @@ -1721,9 +1607,9 @@ gpu_executable { platform_name: "CUDA" cuda_platform_version { cuda_toolkit_version: "12.8.0" - cudnn_version: "9.10.0" + cudnn_version: "9.10.2" cub_version: "3.1.2" } } - buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 64B\n maybe_live_out allocation: 16B\n preallocated temp allocation: 528B\n preallocated temp fragmentation: 0B (0.00%)\n total allocation: 620B\n total fragmentation: 548B (88.39%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 32B\n\t\tXLA Label: fusion\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: async-start\n\t\tShape: ((u32[8]), u32[4])\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 16B\n\t\tXLA Label: reduce-scatter\n\t\tShape: u32[4]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 8B\n\t\tXLA Label: async-start\n\t\tShape: (u32[8])\n\t\t==========================\n\n\tBuffer 7:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 8:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 9:\n\t\tSize: 4B\n\t\tXLA Label: add\n\t\tShape: u32[]\n\t\t==========================\n" + buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 64B\n maybe_live_out allocation: 16B\n preallocated temp allocation: 32B\n preallocated temp fragmentation: 0B (0.00%)\n total allocation: 124B\n total fragmentation: 76B (61.29%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 32B\n\t\tXLA Label: fusion\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: reduce-scatter\n\t\tShape: u32[4]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 7:\n\t\tSize: 4B\n\t\tXLA Label: add\n\t\tShape: u32[]\n\t\t==========================\n" } diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu/v1/AllGatherMixedTypes.pbtxt b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_h100/v2/AllGatherMixedTypes.pbtxt similarity index 60% rename from third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu/v1/AllGatherMixedTypes.pbtxt rename to third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_h100/v2/AllGatherMixedTypes.pbtxt index f75e67bfac744d..ca49a6f758af8b 100644 --- a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu/v1/AllGatherMixedTypes.pbtxt +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_h100/v2/AllGatherMixedTypes.pbtxt @@ -1,9 +1,9 @@ -# proto-file: third_party/tensorflow/compiler/xla/tools/split_proto/split_proto_cli.proto -# proto-message: xla.split_proto_cli.DeserializedSplitExecutableAndOptions - +# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto +# proto-message: xla.HumanReadableAotExecutable executable_and_options { compile_options { executable_build_options { + device_ordinal: -1 result_layout { element_type: TUPLE tuple_shapes { @@ -39,7 +39,7 @@ executable_and_options { xla_cpu_enable_fast_math: false xla_gpu_enable_fast_min_max: false xla_force_host_platform_device_count: 1 - xla_gpu_enable_dynamic_slice_fusion: false + xla_gpu_enable_dynamic_slice_fusion: true xla_hlo_evaluator_use_fast_path: true xla_dump_hlo_as_html: false xla_cpu_fast_math_honor_nans: true @@ -65,7 +65,6 @@ executable_and_options { xla_gpu_use_runtime_fusion: false xla_dump_latency_hiding_schedule: false xla_dump_enable_mlir_pretty_form: true - xla_gpu_enable_latency_hiding_scheduler: false xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP xla_gpu_enable_triton_gemm: true xla_gpu_enable_cudnn_int8x32_convolution_reordering: true @@ -78,14 +77,10 @@ executable_and_options { xla_gpu_all_gather_combine_threshold_bytes: 31457287 xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 xla_gpu_enable_highest_priority_async_stream: true - xla_gpu_enable_pipelined_all_reduce: false xla_gpu_exhaustive_tiling_search: false xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 - xla_gpu_enable_pipelined_all_gather: false xla_gpu_redzone_padding_bytes: 8388608 - xla_gpu_enable_pipelined_reduce_scatter: true - xla_gpu_fused_attention_use_cudnn_rng: false xla_gpu_copy_insertion_use_region_analysis: false xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 xla_gpu_collect_cost_model_stats: false @@ -129,6 +124,7 @@ executable_and_options { xla_reduce_window_rewrite_base_length: 16 xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true xla_gpu_nccl_terminate_on_error: false xla_gpu_shard_autotuning: true xla_gpu_enable_approx_costly_collectives: false @@ -199,7 +195,6 @@ executable_and_options { xla_cpu_generate_unique_c_style_kernel_entry_points: false xla_gpu_experimental_collective_cse_distance_threshold: 0 xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true - xla_cpu_use_fusion_emitters: true xla_gpu_experimental_collective_perf_table_path: "" xla_gpu_collective_permute_combine_threshold_bytes: 31457287 xla_unsupported_crash_on_hlo_pass_noop_change: false @@ -208,7 +203,6 @@ executable_and_options { xla_gpu_unsupported_enable_triton_multi_output_fusion: false xla_gpu_experimental_matmul_perf_table_path: "" xla_gpu_unsupported_enable_all_reduce_decomposer: false - xla_gpu_unsupported_use_all_reduce_one_shot_kernel: true xla_gpu_nccl_blocking_communicators: true xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 xla_gpu_first_collective_call_terminate_timeout_seconds: 40 @@ -219,14 +213,13 @@ executable_and_options { xla_detect_unstable_reductions: DETECTION_MODE_NONE xla_gpu_command_buffer_scheduling_mode: LHS xla_gpu_experimental_enable_nccl_symmetric_buffers: false - xla_gpu_experimental_autotuner_cache_dir: "" xla_disable_automatic_host_compute_offload: false xla_gpu_experimental_enable_fusion_autotuner: true xla_gpu_experimental_scaled_dot_with_triton: true xla_gpu_command_buffer_unroll_loops: false xla_cpu_experimental_onednn_custom_call: false - xla_gpu_experimental_use_raft_select_k: false xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true xla_cpu_collective_call_terminate_timeout_seconds: 40 xla_cpu_collective_call_warn_stuck_seconds: 20 xla_keep_shardings_after_spmd: false @@ -236,31 +229,44 @@ executable_and_options { xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE xla_cpu_enable_platform_dependent_math: true xla_gpu_detect_nan: DETECTION_MODE_NONE - xla_gpu_experimental_dump_gpu_executable: true xla_gpu_detect_inf: DETECTION_MODE_NONE xla_enable_enzyme_comms_opt: false xla_gpu_experimental_enable_buffer_saver_on_thunks: false xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE - xla_gpu_experimental_aot_compiled_thunks: true xla_enable_scoped_logging_timers: true xla_gpu_executable_embed_debug_info: true xla_cpu_collective_timeout_seconds: 1800 xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU xla_gpu_experimental_all_fusions_with_triton: false xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT xla_gpu_rocm_max_trace_events: 4194304 - xla_gpu_experimental_ragged_all_to_all_use_barrier: true xla_gpu_execution_terminate_timeout: "inf" - xla_enable_hlo_sharding_v3: false + xla_enable_hlo_sharding_v3: true xla_gpu_log_minmax: false xla_gpu_print_compilation_stats: false xla_gpu_experimental_enable_conv_fusion: false - xla_gpu_experimental_enable_tiling_propagation: false + xla_gpu_experimental_enable_tiling_propagation: true xla_gpu_execution_progress_tracking: 0 xla_gpu_experimental_max_unroll_factor: 32 xla_gpu_enable_pdl: true xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true xla_recognize_reduction_optimization_level: 0 xla_gpu_experimental_parallel_async_compute_limit: 2 xla_dump_buffer_assignment_analysis: true @@ -280,12 +286,12 @@ executable_and_options { xla_gpu_executable_num_compute_streams: 0 xla_gpu_executable_num_communication_streams: 0 xla_gpu_experimental_collective_start_as_early_as_possible: false - xla_cpu_experimental_enable_tiling_propagation: false + xla_cpu_experimental_enable_tiling_propagation: true xla_gpu_enable_dus_accumulator_zero_init_elimination: false xla_gpu_experimental_thunk_buffer_debug_module_outputs: false xla_cpu_ftz: true xla_gpu_enable_allocator_spatial_partitioning: true - xla_dump_compact_gte: false + xla_dump_compact_gte: true xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES xla_dump_hlo_as_riegeli: false xla_gpu_experimental_use_ragged_dot_grouped_gemm: true @@ -294,6 +300,39 @@ executable_and_options { xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED xla_gpu_enable_nccl_user_buffers_in_default_space: false xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + xla_gpu_enable_persistent_symmetric_memory: false + xla_gpu_experimental_enable_raft_for_stable_topk: false + xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND + xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET } num_replicas: 2 num_partitions: 1 @@ -311,9 +350,9 @@ executable_and_options { allow_spmd_sharding_propagation_to_parameters: false process_count: 1 memory_fitting_level: EFFORT_O2 - slice_size: 2 } } + pjrt_client_name: "PjRtStreamExecutorClient" } gpu_executable { hlo_module_with_config { @@ -338,6 +377,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "param_0.1" } id: 4294967298 frontend_attributes { @@ -361,6 +401,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "bitcast.6.1" } id: 4294967297 operand_ids: 4294967298 @@ -385,6 +426,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "transpose.1.1" } dimensions: 1 dimensions: 0 @@ -444,6 +486,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "param_0.3" } id: 8589934594 frontend_attributes { @@ -467,6 +510,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "bitcast.4.1" } id: 8589934593 operand_ids: 8589934594 @@ -491,6 +535,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "transpose.3" } dimensions: 1 dimensions: 0 @@ -544,6 +589,7 @@ gpu_executable { } } metadata { + scheduling_name: "param_0.5" } id: 12884901890 frontend_attributes { @@ -561,6 +607,7 @@ gpu_executable { } } metadata { + scheduling_name: "p1.1.1" } id: 12884901889 operand_ids: 12884901890 @@ -585,6 +632,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "broadcast.13.1" } id: 12884901888 operand_ids: 12884901889 @@ -618,156 +666,6 @@ gpu_executable { root_id: 12884901888 is_fusion_computation: true } - computations { - name: "async_computation" - instructions { - name: "param_0.6" - opcode: "parameter" - shape { - element_type: U32 - dimensions: 2 - dimensions: 1 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - metadata { - } - id: 17179869184 - frontend_attributes { - } - statistics_viz { - } - } - instructions { - name: "param_1" - opcode: "parameter" - shape { - element_type: F32 - dimensions: 2 - dimensions: 1 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - metadata { - } - parameter_number: 1 - id: 17179869185 - frontend_attributes { - } - statistics_viz { - } - } - instructions { - name: "allgather.2" - opcode: "all-gather" - shape { - element_type: TUPLE - tuple_shapes { - element_type: U32 - dimensions: 2 - dimensions: 2 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - tuple_shapes { - element_type: F32 - dimensions: 2 - dimensions: 2 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - } - metadata { - } - dimensions: 1 - id: 17179869186 - operand_ids: 17179869184 - operand_ids: 17179869185 - frontend_attributes { - } - statistics_viz { - } - collective_device_list { - } - } - program_shape { - parameters { - element_type: U32 - dimensions: 2 - dimensions: 1 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - parameters { - element_type: F32 - dimensions: 2 - dimensions: 1 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - result { - element_type: TUPLE - tuple_shapes { - element_type: U32 - dimensions: 2 - dimensions: 2 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - tuple_shapes { - element_type: F32 - dimensions: 2 - dimensions: 2 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - } - parameter_names: "param_0.6" - parameter_names: "param_1" - } - id: 4 - root_id: 17179869186 - } computations { name: "wrapped_broadcast_computation" instructions { @@ -780,6 +678,7 @@ gpu_executable { } } metadata { + scheduling_name: "param_0.7" } id: 21474836481 frontend_attributes { @@ -803,6 +702,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "broadcast.12.1" } id: 21474836480 operand_ids: 21474836481 @@ -848,6 +748,7 @@ gpu_executable { } } metadata { + scheduling_name: "id.1" } frontend_attributes { } @@ -870,18 +771,16 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "wrapped_broadcast" } fusion_kind: "kLoop" - id: 11 + id: 9 operand_ids: 0 called_computation_ids: 5 frontend_attributes { } statistics_viz { } - backend_config_payload { - id: 0 - } } instructions { name: "loop_broadcast_fusion" @@ -899,6 +798,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "loop_broadcast_fusion" } fusion_kind: "kLoop" id: 8 @@ -908,87 +808,10 @@ gpu_executable { } statistics_viz { } - backend_config_payload { - id: 0 - } - } - instructions { - name: "all-gather-start" - opcode: "async-start" - shape { - element_type: TUPLE - tuple_shapes { - element_type: TUPLE - tuple_shapes { - element_type: U32 - dimensions: 2 - dimensions: 1 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - tuple_shapes { - element_type: F32 - dimensions: 2 - dimensions: 1 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - } - tuple_shapes { - element_type: TUPLE - tuple_shapes { - element_type: U32 - dimensions: 2 - dimensions: 2 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - tuple_shapes { - element_type: F32 - dimensions: 2 - dimensions: 2 - layout { - minor_to_major: 0 - minor_to_major: 1 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - is_dynamic_dimension: false - } - } - } - metadata { - } - id: 9 - operand_ids: 11 - operand_ids: 8 - called_computation_ids: 4 - frontend_attributes { - } - statistics_viz { - } - backend_config_payload { - id: 1 - } } instructions { - name: "all-gather-done" - opcode: "async-done" + name: "allgather.3" + opcode: "all-gather" shape { element_type: TUPLE tuple_shapes { @@ -1017,13 +840,21 @@ gpu_executable { } } metadata { + scheduling_name: "all-gather-start" } + dimensions: 1 id: 10 operand_ids: 9 + operand_ids: 8 frontend_attributes { } statistics_viz { } + collective_device_list { + } + backend_config_payload { + id: 0 + } } instructions { name: "ag0.1" @@ -1041,6 +872,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "ag0.1" } id: 1 operand_ids: 10 @@ -1065,6 +897,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "loop_transpose_fusion.1" } fusion_kind: "kLoop" id: 7 @@ -1074,9 +907,6 @@ gpu_executable { } statistics_viz { } - backend_config_payload { - id: 0 - } } instructions { name: "bitcast.24" @@ -1091,6 +921,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "bitcast.24" } id: 4 operand_ids: 7 @@ -1115,6 +946,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "ag1.1" } tuple_index: 1 id: 2 @@ -1140,6 +972,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "loop_transpose_fusion" } fusion_kind: "kLoop" id: 6 @@ -1149,9 +982,6 @@ gpu_executable { } statistics_viz { } - backend_config_payload { - id: 0 - } } instructions { name: "bitcast.1.0" @@ -1166,6 +996,7 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "bitcast.1.0" } id: 5 operand_ids: 6 @@ -1199,6 +1030,7 @@ gpu_executable { } } metadata { + scheduling_name: "out.1" } id: 3 operand_ids: 4 @@ -1262,26 +1094,17 @@ gpu_executable { value { instruction_ids: 0 instruction_ids: 8 - instruction_ids: 11 instruction_ids: 9 instruction_ids: 10 - instruction_ids: 2 instruction_ids: 1 instruction_ids: 7 - instruction_ids: 4 + instruction_ids: 2 instruction_ids: 6 + instruction_ids: 4 instruction_ids: 5 instruction_ids: 3 } } - sequences { - key: 4 - value { - instruction_ids: 17179869185 - instruction_ids: 17179869184 - instruction_ids: 17179869186 - } - } } input_output_alias { } @@ -1298,11 +1121,10 @@ gpu_executable { frontend_attributes { map { key: "fingerprint_before_lhs" - value: "27beb9c461cc00eca8e7c7d6d313ca37" + value: "7b2e4f79c5524b5cbc566da178aa7fda" } } - payloads: "{\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"native_emitter_backend_config\":{\"type\":\"NATIVE_EMITTER_TYPE_INVALID\",\"unroll_factor\":0},\"operation_queue_id\":\"0\",\"reification_cost\":[]}" - payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" + payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" } config { entry_computation_layout { @@ -1345,7 +1167,7 @@ gpu_executable { xla_cpu_enable_fast_math: false xla_gpu_enable_fast_min_max: false xla_force_host_platform_device_count: 1 - xla_gpu_enable_dynamic_slice_fusion: false + xla_gpu_enable_dynamic_slice_fusion: true xla_hlo_evaluator_use_fast_path: true xla_dump_hlo_as_html: false xla_cpu_fast_math_honor_nans: true @@ -1371,7 +1193,6 @@ gpu_executable { xla_gpu_use_runtime_fusion: false xla_dump_latency_hiding_schedule: false xla_dump_enable_mlir_pretty_form: true - xla_gpu_enable_latency_hiding_scheduler: false xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP xla_gpu_enable_triton_gemm: true xla_gpu_enable_cudnn_int8x32_convolution_reordering: true @@ -1384,14 +1205,10 @@ gpu_executable { xla_gpu_all_gather_combine_threshold_bytes: 31457287 xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 xla_gpu_enable_highest_priority_async_stream: true - xla_gpu_enable_pipelined_all_reduce: false xla_gpu_exhaustive_tiling_search: false xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 - xla_gpu_enable_pipelined_all_gather: false xla_gpu_redzone_padding_bytes: 8388608 - xla_gpu_enable_pipelined_reduce_scatter: true - xla_gpu_fused_attention_use_cudnn_rng: false xla_gpu_copy_insertion_use_region_analysis: false xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 xla_gpu_collect_cost_model_stats: false @@ -1435,6 +1252,7 @@ gpu_executable { xla_reduce_window_rewrite_base_length: 16 xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true xla_gpu_nccl_terminate_on_error: false xla_gpu_shard_autotuning: true xla_gpu_enable_approx_costly_collectives: false @@ -1505,7 +1323,6 @@ gpu_executable { xla_cpu_generate_unique_c_style_kernel_entry_points: false xla_gpu_experimental_collective_cse_distance_threshold: 0 xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true - xla_cpu_use_fusion_emitters: true xla_gpu_experimental_collective_perf_table_path: "" xla_gpu_collective_permute_combine_threshold_bytes: 31457287 xla_unsupported_crash_on_hlo_pass_noop_change: false @@ -1514,7 +1331,6 @@ gpu_executable { xla_gpu_unsupported_enable_triton_multi_output_fusion: false xla_gpu_experimental_matmul_perf_table_path: "" xla_gpu_unsupported_enable_all_reduce_decomposer: false - xla_gpu_unsupported_use_all_reduce_one_shot_kernel: true xla_gpu_nccl_blocking_communicators: true xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 xla_gpu_first_collective_call_terminate_timeout_seconds: 40 @@ -1525,14 +1341,13 @@ gpu_executable { xla_detect_unstable_reductions: DETECTION_MODE_NONE xla_gpu_command_buffer_scheduling_mode: LHS xla_gpu_experimental_enable_nccl_symmetric_buffers: false - xla_gpu_experimental_autotuner_cache_dir: "" xla_disable_automatic_host_compute_offload: false xla_gpu_experimental_enable_fusion_autotuner: true xla_gpu_experimental_scaled_dot_with_triton: true xla_gpu_command_buffer_unroll_loops: false xla_cpu_experimental_onednn_custom_call: false - xla_gpu_experimental_use_raft_select_k: false xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true xla_cpu_collective_call_terminate_timeout_seconds: 40 xla_cpu_collective_call_warn_stuck_seconds: 20 xla_keep_shardings_after_spmd: false @@ -1542,31 +1357,44 @@ gpu_executable { xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE xla_cpu_enable_platform_dependent_math: true xla_gpu_detect_nan: DETECTION_MODE_NONE - xla_gpu_experimental_dump_gpu_executable: true xla_gpu_detect_inf: DETECTION_MODE_NONE xla_enable_enzyme_comms_opt: false xla_gpu_experimental_enable_buffer_saver_on_thunks: false xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE - xla_gpu_experimental_aot_compiled_thunks: true xla_enable_scoped_logging_timers: true xla_gpu_executable_embed_debug_info: true xla_cpu_collective_timeout_seconds: 1800 xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU xla_gpu_experimental_all_fusions_with_triton: false xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT xla_gpu_rocm_max_trace_events: 4194304 - xla_gpu_experimental_ragged_all_to_all_use_barrier: true xla_gpu_execution_terminate_timeout: "inf" - xla_enable_hlo_sharding_v3: false + xla_enable_hlo_sharding_v3: true xla_gpu_log_minmax: false xla_gpu_print_compilation_stats: false xla_gpu_experimental_enable_conv_fusion: false - xla_gpu_experimental_enable_tiling_propagation: false + xla_gpu_experimental_enable_tiling_propagation: true xla_gpu_execution_progress_tracking: 0 xla_gpu_experimental_max_unroll_factor: 32 xla_gpu_enable_pdl: true xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true xla_recognize_reduction_optimization_level: 0 xla_gpu_experimental_parallel_async_compute_limit: 2 xla_dump_buffer_assignment_analysis: true @@ -1586,12 +1414,12 @@ gpu_executable { xla_gpu_executable_num_compute_streams: 0 xla_gpu_executable_num_communication_streams: 0 xla_gpu_experimental_collective_start_as_early_as_possible: false - xla_cpu_experimental_enable_tiling_propagation: false + xla_cpu_experimental_enable_tiling_propagation: true xla_gpu_enable_dus_accumulator_zero_init_elimination: false xla_gpu_experimental_thunk_buffer_debug_module_outputs: false xla_cpu_ftz: true xla_gpu_enable_allocator_spatial_partitioning: true - xla_dump_compact_gte: false + xla_dump_compact_gte: true xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES xla_dump_hlo_as_riegeli: false xla_gpu_experimental_use_ragged_dot_grouped_gemm: true @@ -1600,6 +1428,39 @@ gpu_executable { xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED xla_gpu_enable_nccl_user_buffers_in_default_space: false xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + xla_gpu_enable_persistent_symmetric_memory: false + xla_gpu_experimental_enable_raft_for_stable_topk: false + xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND + xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET } static_device_assignment { replica_count: 2 @@ -1620,95 +1481,65 @@ gpu_executable { } } buffer_assignment { - logical_buffers { - id: 7 - size: 16 - defined_at { - instruction_id: 17179869186 - } - } - logical_buffers { - id: 8 - size: 16 - defined_at { - shape_index: 0 - instruction_id: 17179869186 - } - } logical_buffers { id: 9 - size: 16 - defined_at { - shape_index: 1 - instruction_id: 17179869186 - } - } - logical_buffers { - id: 12 size: 4 defined_at { } } logical_buffers { - id: 13 + id: 10 size: 8 defined_at { - instruction_id: 11 + instruction_id: 9 } } logical_buffers { - id: 14 + id: 11 size: 8 defined_at { instruction_id: 8 } } logical_buffers { - id: 15 + id: 12 size: 16 defined_at { - instruction_id: 9 + instruction_id: 10 } } logical_buffers { - id: 16 + id: 13 size: 16 defined_at { shape_index: 0 - instruction_id: 9 + instruction_id: 10 } } logical_buffers { - id: 17 + id: 14 size: 16 defined_at { shape_index: 1 - instruction_id: 9 - } - } - logical_buffers { - id: 18 - size: 16 - defined_at { instruction_id: 10 } } logical_buffers { - id: 19 + id: 15 size: 16 defined_at { instruction_id: 7 } } logical_buffers { - id: 20 + id: 16 size: 16 defined_at { instruction_id: 6 } } logical_buffers { - id: 21 + id: 17 size: 16 defined_at { instruction_id: 3 @@ -1718,12 +1549,12 @@ gpu_executable { size: 16 maybe_live_out: true assigned { - logical_buffer_id: 13 + logical_buffer_id: 10 size: 8 element_type: U32 } assigned { - logical_buffer_id: 19 + logical_buffer_id: 15 size: 16 element_type: U32 } @@ -1733,17 +1564,17 @@ gpu_executable { size: 16 maybe_live_out: true assigned { - logical_buffer_id: 8 - size: 16 + logical_buffer_id: 9 + size: 4 element_type: U32 } assigned { - logical_buffer_id: 12 - size: 4 + logical_buffer_id: 13 + size: 16 element_type: U32 } assigned { - logical_buffer_id: 20 + logical_buffer_id: 16 size: 16 element_type: F32 } @@ -1753,7 +1584,7 @@ gpu_executable { size: 16 maybe_live_out: true assigned { - logical_buffer_id: 21 + logical_buffer_id: 17 size: 16 element_type: TUPLE } @@ -1761,102 +1592,74 @@ gpu_executable { } buffer_allocations { index: 3 - size: 1544 + size: 520 assigned { - logical_buffer_id: 7 - offset: 1024 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 9 - offset: 1280 - size: 16 - element_type: F32 - } - assigned { - logical_buffer_id: 14 - offset: 1536 + logical_buffer_id: 11 + offset: 512 size: 8 element_type: F32 } assigned { - logical_buffer_id: 15 + logical_buffer_id: 12 size: 16 element_type: TUPLE } assigned { - logical_buffer_id: 16 + logical_buffer_id: 14 offset: 256 size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 17 - offset: 512 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 18 - offset: 768 - size: 16 - element_type: TUPLE + element_type: F32 } is_tuple: true } heap_simulator_traces { events { - buffer_id: 14 + buffer_id: 11 computation_name: "test_computation" instruction_name: "loop_broadcast_fusion" } events { - buffer_id: 9 - computation_name: "async_computation" - instruction_name: "allgather.2" + buffer_id: 14 + computation_name: "test_computation" + instruction_name: "allgather.3" } events { kind: FREE - buffer_id: 14 + buffer_id: 11 computation_name: "test_computation" instruction_name: "loop_broadcast_fusion" } events { kind: FREE - buffer_id: 9 - computation_name: "async_computation" - instruction_name: "allgather.2" + buffer_id: 14 + computation_name: "test_computation" + instruction_name: "allgather.3" } whole_module_simulation: true buffer_allocation_index: 3 } peak_buffers { - logical_buffer_ids: 19 + logical_buffer_ids: 15 } peak_buffers { index: 1 - logical_buffer_ids: 20 + logical_buffer_ids: 16 } peak_buffers { index: 2 - logical_buffer_ids: 21 + logical_buffer_ids: 17 } peak_buffers { index: 3 - logical_buffer_ids: 15 - logical_buffer_ids: 16 - logical_buffer_ids: 17 - logical_buffer_ids: 18 - logical_buffer_ids: 7 - logical_buffer_ids: 9 + logical_buffer_ids: 12 + logical_buffer_ids: 11 logical_buffer_ids: 14 } } gpu_compute_capability { cuda_compute_capability { - major: 6 - feature_extension: NONE + major: 9 + feature_extension: ACCELERATED_FEATURES } } module_name: "test" @@ -1887,92 +1690,21 @@ gpu_executable { values { size: 16 maybe_live_out: true - assigned { - logical_buffer_id: 13 - size: 8 - element_type: U32 - } - assigned { - logical_buffer_id: 19 - size: 16 - element_type: U32 - } } values { index: 1 size: 16 maybe_live_out: true - assigned { - logical_buffer_id: 8 - size: 16 - element_type: U32 - } - assigned { - logical_buffer_id: 12 - size: 4 - element_type: U32 - } - assigned { - logical_buffer_id: 20 - size: 16 - element_type: F32 - } } values { index: 2 size: 16 maybe_live_out: true - assigned { - logical_buffer_id: 21 - size: 16 - element_type: TUPLE - } is_tuple: true } values { index: 3 - size: 1544 - assigned { - logical_buffer_id: 7 - offset: 1024 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 9 - offset: 1280 - size: 16 - element_type: F32 - } - assigned { - logical_buffer_id: 14 - offset: 1536 - size: 8 - element_type: F32 - } - assigned { - logical_buffer_id: 15 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 16 - offset: 256 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 17 - offset: 512 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 18 - offset: 768 - size: 16 - element_type: TUPLE - } + size: 520 is_tuple: true } } @@ -2032,7 +1764,7 @@ gpu_executable { } args { slice { - offset: 1536 + offset: 512 size: 8 buffer_allocation_index: 3 element_type: F32 @@ -2056,7 +1788,7 @@ gpu_executable { name: "loop_broadcast_fusion" kernel_spec { cubin { - data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000@\010\000\000\000\000\000\000\300\005\000\000\000\000\000\000<\005<\000@\0008\000\003\000@\000\n\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.constant0.loop_broadcast_fusion\000.rel.nv.constant0.loop_broadcast_fusion\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.rel.nv.constant0.loop_broadcast_fusion\000.nv.constant0.loop_broadcast_fusion\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000loop_broadcast_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\266\000\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\332\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\000\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\300\000\000\000\000\000\000\000\004/\010\000\005\000\000\000\007\000\000\000\004\021\010\000\005\000\000\000\000\000\000\000\004\022\010\000\005\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\0010\000\000\001*\000\000\004\n\010\000\002\000\000\000@\001\020\000\003\031\020\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003\033\377\000\004\034\004\000x\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377s\000\000\000\000\000\000\000\000\000\000\021%\000\0056\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\007\000\376\000D\034\000\001\000\207\000\200\007\230L\002\000\007\005\200\007\230L\004\000\027\002\000\000\310\360\342\007\340\367\000\304?\010\003\000\027\005\200\007\230L\002\002\007\000\000\240\324\356\000\004\347\001\000\000(8\346\007\340\374\000t\034\004\004\004\'\005\000\201\030L\005\0007\005\000\010\020L\006\n\'\000\000\000\270\\\361\017\200\374\000\374\037\000\006\004\007\000\000 \334\356\000\017\007\000\000\000\260P\017\000\007\000\000\000\000\343\377\007\000\374\000\200\037\000\017\000\007\377\377\017@\342\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\340\007\000\374\000\200\037\000\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\005\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000E\001\000\000\000\000\000\000\033\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000`\002\000\000\000\000\000\000\220\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\002\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000N\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\003\000\000\000\000\000\000\\\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\332\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\003\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\366\000\000\000\013\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\216\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\003\000\000\000\000\000\000P\001\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\300\000\000\000\000\000\000\000\003\000\000\000\005\000\000\007 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\005\000\000\000@\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\240\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000@\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\240\013\000\000\000\000\000\000\240\010\000\000\000\000\000\000Z\rZ\000@\0008\000\005\000@\000\014\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_broadcast_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_broadcast_fusion\000.nv.constant0.loop_broadcast_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\216\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\275\000\000\000 \240\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\336\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000*\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000@\001\000\000\003\000\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004(\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\004\034\004\000\240\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\020\002\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\000\000\000\000D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\n\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\204\000\000\000\n\000\000\000\"\016\000\271z\004\000\000\202\000\000\000\n\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\206\000\000\000\n\000\000\000b\016\000\201y\002\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\342/\000Er\007\000\002\000\000\000\000\020 \000\000\312O\000\206y\000\004\007\000\000\000\004\031\020\014\000\342\017\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\024\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\211\001\000\000\000\000\000\000d\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\002\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\244\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\003\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\004\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000N\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000<\004\000\000\000\000\000\000`\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\324\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\234\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\302\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\004\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\216\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000 \002\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000" } arity: 2 kernel_name: "loop_broadcast_fusion" @@ -2089,6 +1821,7 @@ gpu_executable { } } } + use_pdl: true tma_metadata { } } @@ -2136,7 +1869,7 @@ gpu_executable { name: "wrapped_broadcast" kernel_spec { cubin { - data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\340\007\000\000\000\000\000\000`\005\000\000\000\000\000\000<\005<\000@\0008\000\003\000@\000\n\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.constant0.wrapped_broadcast\000.rel.nv.constant0.wrapped_broadcast\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.rel.nv.constant0.wrapped_broadcast\000.nv.constant0.wrapped_broadcast\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000wrapped_broadcast\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\246\000\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\306\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\342\000\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\361\000\000\000\022\020\t\000\000\000\000\000\000\000\000\000\200\000\000\000\000\000\000\000\004/\010\000\005\000\000\000\006\000\000\000\004\021\010\000\005\000\000\000\000\000\000\000\004\022\010\000\005\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\0010\000\000\001*\000\000\004\n\010\000\002\000\000\000@\001\020\000\003\031\020\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003\033\377\000\004\034\004\000p\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377s\000\000\000\000\000\000\000\000\000\000\021%\000\0056\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\007\000\376\000D\034\000\001\000\207\000\200\007\230L\002\000\007\005\200\007\230L\004\000\027\002\000\000\310\360\342\007\340\367\000\304?\010\003\000\027\005\200\007\230L\002\002\007\000\000\240\324\356\000\004\347\001\000\000(8\346\007\340\374\000\304\037\004\004\004\'\005\000\201\030L\005\0007\005\000\010\020L\002\004\007\000\000 \334\356\352\007\340\377\000\374\037\000\000\017\007\000\000\000\260P\017\000\007\000\000\000\000\343\017\000\207\377\377\017@\342\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\361\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0001\001\000\000\000\000\000\000\003\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\002\000\000\000\000\000\000\220\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\002\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000J\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\354\002\000\000\000\000\000\000\\\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\306\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\003\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\342\000\000\000\013\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000h\003\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\202\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\003\000\000\000\000\000\000P\001\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\004\000\000\000\000\000\000\200\000\000\000\000\000\000\000\003\000\000\000\005\000\000\006 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\005\000\000\000\340\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000x\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\350\001\000\000\000\000\000\000\350\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\340\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\240\013\000\000\000\000\000\000\240\010\000\000\000\000\000\000Z\rZ\000@\0008\000\005\000@\000\014\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.wrapped_broadcast\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000wrapped_broadcast\000.nv.constant0.wrapped_broadcast\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\261\000\000\000 \240\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\322\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\001\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\036\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\0000\001\000\000\003\000\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004$\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\004\034\004\000\220\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\020\002\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\000\000\000\000D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\n\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\204\000\000\000\n\000\000\000\"\016\000\271z\004\000\000\202\000\000\000\n\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\206\000\000\000\n\000\000\000b\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\312/\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\004\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000u\001\000\000\000\000\000\000P\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\002\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\230\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\003\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\003\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000J\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\004\000\000\000\000\000\000`\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000t\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\266\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\230\004\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\344\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000 \002\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000" } arity: 2 kernel_name: "wrapped_broadcast" @@ -2169,13 +1902,14 @@ gpu_executable { } } } + use_pdl: true tma_metadata { } } } thunks { thunk_info { - profile_annotation: "allgather.2" + profile_annotation: "allgather.3" thunk_id: 4 } all_gather_thunk { @@ -2227,7 +1961,7 @@ gpu_executable { element_count: 2 source_buffer { slice { - offset: 1536 + offset: 512 size: 8 buffer_allocation_index: 3 element_type: F32 @@ -2247,7 +1981,7 @@ gpu_executable { } destination_buffer { slice { - offset: 1280 + offset: 256 size: 16 buffer_allocation_index: 3 element_type: F32 @@ -2317,7 +2051,7 @@ gpu_executable { name: "loop_transpose_fusion_1" kernel_spec { cubin { - data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000`\010\000\000\000\000\000\000\340\005\000\000\000\000\000\000<\005<\000@\0008\000\003\000@\000\n\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.constant0.loop_transpose_fusion_1\000.rel.nv.constant0.loop_transpose_fusion_1\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.rel.nv.constant0.loop_transpose_fusion_1\000.nv.constant0.loop_transpose_fusion_1\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000loop_transpose_fusion_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\276\000\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\344\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\017\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\300\000\000\000\000\000\000\000\004/\010\000\005\000\000\000\006\000\000\000\004\021\010\000\005\000\000\000\000\000\000\000\004\022\010\000\005\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\0010\000\000\001*\000\000\004\n\010\000\002\000\000\000@\001\020\000\003\031\020\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003\033\377\000\004\034\004\000\220\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377s\000\000\000\000\000\000\000\000\000\000\021%\000\0056\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\007\340\343\000\230?\000\001\000\207\000\200\007\230L\005\000\027\002\000\000\310\360\000\005W\000\000\000\020\\\366\007\300\376\000\304\037\010\000\000\'\000\000\000\000\004\000\005\007\360\201\177\337[\003\000\347\001\000\000(8\346\007@\374\000\304\036\000\002\000\007\005\000\201\030L\003\003\027\005\000\010\020L\002\002\007\000\000\240\324\356\361\007\302\374\000\234\037\000\000\005\347\001\000\000(8\004\005\'\005\000\201\030L\005\0007\005\000\010\020L\361\007\341\377\000\374\037\000\002\004\007\000\000 \334\356\017\000\007\000\000\000\000\343\017\000\207\377\377\017@\342\340\007\000\374\000\200\037\000\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\017\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000O\001\000\000\000\000\000\000\'\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\002\000\000\000\000\000\000\220\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\003\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000P\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000,\003\000\000\000\000\000\000\\\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\344\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\003\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\001\000\000\013\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\003\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\224\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\270\003\000\000\000\000\000\000P\001\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \005\000\000\000\000\000\000\300\000\000\000\000\000\000\000\003\000\000\000\005\000\000\006 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\005\000\000\000`\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\270\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000(\002\000\000\000\000\000\000(\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000`\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\240\013\000\000\000\000\000\000\240\010\000\000\000\000\000\000Z\rZ\000@\0008\000\005\000@\000\014\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion_1\000.nv.constant0.loop_transpose_fusion_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\224\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\303\000\000\000 \240\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\344\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\001\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0000\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000H\001\000\000\003\000\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\020\002\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\000\000\000\000D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\n\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\204\000\000\000\n\000\000\000b\016\000\271z\004\000\000\202\000\000\000\n\000\000\000\342\017\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\206\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242\016\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\034\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\223\001\000\000\000\000\000\000n\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\003\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\252\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\003\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0000\004\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000P\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\004\000\000\000\000\000\000`\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\332\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\264\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\310\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\330\004\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\224\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000 \002\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000" } arity: 2 kernel_name: "loop_transpose_fusion_1" @@ -2350,6 +2084,7 @@ gpu_executable { } } } + use_pdl: true tma_metadata { } } @@ -2362,7 +2097,7 @@ gpu_executable { custom_kernel_thunk { args { slice { - offset: 1280 + offset: 256 size: 16 buffer_allocation_index: 3 element_type: F32 @@ -2405,7 +2140,7 @@ gpu_executable { name: "loop_transpose_fusion" kernel_spec { cubin { - data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000@\010\000\000\000\000\000\000\300\005\000\000\000\000\000\000<\005<\000@\0008\000\003\000@\000\n\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.constant0.loop_transpose_fusion\000.rel.nv.constant0.loop_transpose_fusion\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.rel.nv.constant0.loop_transpose_fusion\000.nv.constant0.loop_transpose_fusion\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000loop_transpose_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\266\000\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\332\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\000\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\300\000\000\000\000\000\000\000\004/\010\000\005\000\000\000\006\000\000\000\004\021\010\000\005\000\000\000\000\000\000\000\004\022\010\000\005\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\0010\000\000\001*\000\000\004\n\010\000\002\000\000\000@\001\020\000\003\031\020\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003\033\377\000\004\034\004\000\220\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377s\000\000\000\000\000\000\000\000\000\000\021%\000\0056\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\007\340\343\000\230?\000\001\000\207\000\200\007\230L\005\000\027\002\000\000\310\360\000\005W\000\000\000\020\\\366\007\300\376\000\304\037\010\000\000\'\000\000\000\000\004\000\005\007\360\201\177\337[\003\000\347\001\000\000(8\346\007@\374\000\304\036\000\002\000\007\005\000\201\030L\003\003\027\005\000\010\020L\002\002\007\000\000\240\324\356\361\007\302\374\000\234\037\000\000\005\347\001\000\000(8\004\005\'\005\000\201\030L\005\0007\005\000\010\020L\361\007\341\377\000\374\037\000\002\004\007\000\000 \334\356\017\000\007\000\000\000\000\343\017\000\207\377\377\017@\342\340\007\000\374\000\200\037\000\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\005\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000E\001\000\000\000\000\000\000\033\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000`\002\000\000\000\000\000\000\220\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\002\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000N\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\024\003\000\000\000\000\000\000\\\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\332\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\003\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\366\000\000\000\013\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\216\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\003\000\000\000\000\000\000P\001\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\300\000\000\000\000\000\000\000\003\000\000\000\005\000\000\006 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\005\000\000\000@\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\240\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000@\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\240\013\000\000\000\000\000\000\240\010\000\000\000\000\000\000Z\rZ\000@\0008\000\005\000@\000\014\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion\000.nv.constant0.loop_transpose_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\216\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\275\000\000\000 \240\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\336\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\016\001\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000*\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000@\001\000\000\003\000\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\004/\010\000\006\000\000\000\n\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\007\000\000\000\020\002\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\000\000\000\000D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\n\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\204\000\000\000\n\000\000\000b\016\000\271z\004\000\000\202\000\000\000\n\000\000\000\342\017\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\206\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242\016\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\024\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\211\001\000\000\000\000\000\000d\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\002\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\244\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\003\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\004\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000N\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000<\004\000\000\000\000\000\000`\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\324\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\234\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\302\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\004\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\216\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000 \002\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\240\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000" } arity: 2 kernel_name: "loop_transpose_fusion" @@ -2438,6 +2173,7 @@ gpu_executable { } } } + use_pdl: true tma_metadata { } } @@ -2446,9 +2182,9 @@ gpu_executable { platform_name: "CUDA" cuda_platform_version { cuda_toolkit_version: "12.8.0" - cudnn_version: "9.10.0" + cudnn_version: "9.10.2" cub_version: "3.1.2" } } - buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 0B\n maybe_live_out allocation: 48B\n preallocated temp allocation: 1.5KiB\n preallocated temp fragmentation: 240B (15.54%)\n total allocation: 1.6KiB\n total fragmentation: 1.4KiB (91.96%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: (u32[2,2], f32[2,2])\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 16B\n\t\tXLA Label: async-done\n\t\tShape: (u32[2,2], f32[2,2])\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: async-start\n\t\tShape: (u32[2,2], f32[2,2])\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 16B\n\t\tXLA Label: async-start\n\t\tShape: (u32[2,1], f32[2,1])\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 16B\n\t\tXLA Label: async-start\n\t\tShape: ((u32[2,1], f32[2,1]), (u32[2,2], f32[2,2]))\n\t\t==========================\n\n\tBuffer 7:\n\t\tSize: 16B\n\t\tXLA Label: tuple\n\t\tShape: (u32[4], f32[4])\n\t\t==========================\n\n\tBuffer 8:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 9:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: u32[2,2]\n\t\t==========================\n\n\tBuffer 10:\n\t\tSize: 8B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,1]\n\t\t==========================\n" + buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 0B\n maybe_live_out allocation: 48B\n preallocated temp allocation: 520B\n preallocated temp fragmentation: 240B (46.15%)\n total allocation: 568B\n total fragmentation: 504B (88.73%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: (u32[2,2], f32[2,2])\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 16B\n\t\tXLA Label: tuple\n\t\tShape: (u32[4], f32[4])\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: u32[2,2]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 8B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,1]\n\t\t==========================\n" } diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu/v1/ReduceScatter.pbtxt b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_h100/v2/ReduceScatter.pbtxt similarity index 66% rename from third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu/v1/ReduceScatter.pbtxt rename to third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_h100/v2/ReduceScatter.pbtxt index dbf790a32fd03a..fcd7a336a45f1c 100644 --- a/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu/v1/ReduceScatter.pbtxt +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/gpu/executables/collective_ops_aot_test_2gpu_h100/v2/ReduceScatter.pbtxt @@ -1,9 +1,9 @@ -# proto-file: third_party/tensorflow/compiler/xla/tools/split_proto/split_proto_cli.proto -# proto-message: xla.split_proto_cli.DeserializedSplitExecutableAndOptions - +# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto +# proto-message: xla.HumanReadableAotExecutable executable_and_options { compile_options { executable_build_options { + device_ordinal: -1 result_layout { element_type: U32 dimensions: 4 @@ -27,7 +27,7 @@ executable_and_options { xla_cpu_enable_fast_math: false xla_gpu_enable_fast_min_max: false xla_force_host_platform_device_count: 1 - xla_gpu_enable_dynamic_slice_fusion: false + xla_gpu_enable_dynamic_slice_fusion: true xla_hlo_evaluator_use_fast_path: true xla_dump_hlo_as_html: false xla_cpu_fast_math_honor_nans: true @@ -53,7 +53,6 @@ executable_and_options { xla_gpu_use_runtime_fusion: false xla_dump_latency_hiding_schedule: false xla_dump_enable_mlir_pretty_form: true - xla_gpu_enable_latency_hiding_scheduler: false xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP xla_gpu_enable_triton_gemm: true xla_gpu_enable_cudnn_int8x32_convolution_reordering: true @@ -66,14 +65,10 @@ executable_and_options { xla_gpu_all_gather_combine_threshold_bytes: 31457287 xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 xla_gpu_enable_highest_priority_async_stream: true - xla_gpu_enable_pipelined_all_reduce: false xla_gpu_exhaustive_tiling_search: false xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 - xla_gpu_enable_pipelined_all_gather: false xla_gpu_redzone_padding_bytes: 8388608 - xla_gpu_enable_pipelined_reduce_scatter: true - xla_gpu_fused_attention_use_cudnn_rng: false xla_gpu_copy_insertion_use_region_analysis: false xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 xla_gpu_collect_cost_model_stats: false @@ -117,6 +112,7 @@ executable_and_options { xla_reduce_window_rewrite_base_length: 16 xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true xla_gpu_nccl_terminate_on_error: false xla_gpu_shard_autotuning: true xla_gpu_enable_approx_costly_collectives: false @@ -187,7 +183,6 @@ executable_and_options { xla_cpu_generate_unique_c_style_kernel_entry_points: false xla_gpu_experimental_collective_cse_distance_threshold: 0 xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true - xla_cpu_use_fusion_emitters: true xla_gpu_experimental_collective_perf_table_path: "" xla_gpu_collective_permute_combine_threshold_bytes: 31457287 xla_unsupported_crash_on_hlo_pass_noop_change: false @@ -196,7 +191,6 @@ executable_and_options { xla_gpu_unsupported_enable_triton_multi_output_fusion: false xla_gpu_experimental_matmul_perf_table_path: "" xla_gpu_unsupported_enable_all_reduce_decomposer: false - xla_gpu_unsupported_use_all_reduce_one_shot_kernel: true xla_gpu_nccl_blocking_communicators: true xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 xla_gpu_first_collective_call_terminate_timeout_seconds: 40 @@ -207,14 +201,13 @@ executable_and_options { xla_detect_unstable_reductions: DETECTION_MODE_NONE xla_gpu_command_buffer_scheduling_mode: LHS xla_gpu_experimental_enable_nccl_symmetric_buffers: false - xla_gpu_experimental_autotuner_cache_dir: "" xla_disable_automatic_host_compute_offload: false xla_gpu_experimental_enable_fusion_autotuner: true xla_gpu_experimental_scaled_dot_with_triton: true xla_gpu_command_buffer_unroll_loops: false xla_cpu_experimental_onednn_custom_call: false - xla_gpu_experimental_use_raft_select_k: false xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true xla_cpu_collective_call_terminate_timeout_seconds: 40 xla_cpu_collective_call_warn_stuck_seconds: 20 xla_keep_shardings_after_spmd: false @@ -224,31 +217,44 @@ executable_and_options { xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE xla_cpu_enable_platform_dependent_math: true xla_gpu_detect_nan: DETECTION_MODE_NONE - xla_gpu_experimental_dump_gpu_executable: true xla_gpu_detect_inf: DETECTION_MODE_NONE xla_enable_enzyme_comms_opt: false xla_gpu_experimental_enable_buffer_saver_on_thunks: false xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE - xla_gpu_experimental_aot_compiled_thunks: true xla_enable_scoped_logging_timers: true xla_gpu_executable_embed_debug_info: true xla_cpu_collective_timeout_seconds: 1800 xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU xla_gpu_experimental_all_fusions_with_triton: false xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT xla_gpu_rocm_max_trace_events: 4194304 - xla_gpu_experimental_ragged_all_to_all_use_barrier: true xla_gpu_execution_terminate_timeout: "inf" - xla_enable_hlo_sharding_v3: false + xla_enable_hlo_sharding_v3: true xla_gpu_log_minmax: false xla_gpu_print_compilation_stats: false xla_gpu_experimental_enable_conv_fusion: false - xla_gpu_experimental_enable_tiling_propagation: false + xla_gpu_experimental_enable_tiling_propagation: true xla_gpu_execution_progress_tracking: 0 xla_gpu_experimental_max_unroll_factor: 32 xla_gpu_enable_pdl: true xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true xla_recognize_reduction_optimization_level: 0 xla_gpu_experimental_parallel_async_compute_limit: 2 xla_dump_buffer_assignment_analysis: true @@ -268,12 +274,12 @@ executable_and_options { xla_gpu_executable_num_compute_streams: 0 xla_gpu_executable_num_communication_streams: 0 xla_gpu_experimental_collective_start_as_early_as_possible: false - xla_cpu_experimental_enable_tiling_propagation: false + xla_cpu_experimental_enable_tiling_propagation: true xla_gpu_enable_dus_accumulator_zero_init_elimination: false xla_gpu_experimental_thunk_buffer_debug_module_outputs: false xla_cpu_ftz: true xla_gpu_enable_allocator_spatial_partitioning: true - xla_dump_compact_gte: false + xla_dump_compact_gte: true xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES xla_dump_hlo_as_riegeli: false xla_gpu_experimental_use_ragged_dot_grouped_gemm: true @@ -282,6 +288,39 @@ executable_and_options { xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED xla_gpu_enable_nccl_user_buffers_in_default_space: false xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + xla_gpu_enable_persistent_symmetric_memory: false + xla_gpu_experimental_enable_raft_for_stable_topk: false + xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND + xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET } num_replicas: 2 num_partitions: 1 @@ -299,9 +338,9 @@ executable_and_options { allow_spmd_sharding_propagation_to_parameters: false process_count: 1 memory_fitting_level: EFFORT_O2 - slice_size: 2 } } + pjrt_client_name: "PjRtStreamExecutorClient" } gpu_executable { hlo_module_with_config { @@ -320,6 +359,7 @@ gpu_executable { } } metadata { + scheduling_name: "lhs" } frontend_attributes { } @@ -336,6 +376,7 @@ gpu_executable { } } metadata { + scheduling_name: "rhs" } parameter_number: 1 id: 1 @@ -354,6 +395,7 @@ gpu_executable { } } metadata { + scheduling_name: "add.1" } id: 2 operand_ids: 0 @@ -399,9 +441,10 @@ gpu_executable { } } metadata { + scheduling_name: "param_2.2" } parameter_number: 2 - id: 12884901893 + id: 8589934597 frontend_attributes { } statistics_viz { @@ -427,14 +470,14 @@ gpu_executable { } u32s: 0 } - id: 12884901894 + id: 8589934598 frontend_attributes { } statistics_viz { } } instructions { - name: "p.4" + name: "p.2" opcode: "compare" shape { element_type: PRED @@ -443,19 +486,20 @@ gpu_executable { } } metadata { + scheduling_name: "p.2" } - id: 12884901892 - operand_ids: 12884901893 - operand_ids: 12884901894 + id: 8589934596 + operand_ids: 8589934597 + operand_ids: 8589934598 comparison_direction: "EQ" frontend_attributes { } - comparison_type: "UNSIGNED" statistics_viz { } + comparison_order: "TOTAL" } instructions { - name: "pb.4" + name: "pb.2" opcode: "broadcast" shape { element_type: PRED @@ -467,16 +511,17 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "pb.2" } - id: 12884901889 - operand_ids: 12884901892 + id: 8589934593 + operand_ids: 8589934596 frontend_attributes { } statistics_viz { } } instructions { - name: "param_0.4" + name: "param_0.1" opcode: "parameter" shape { element_type: U32 @@ -488,15 +533,16 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "param_0.1" } - id: 12884901890 + id: 8589934594 frontend_attributes { } statistics_viz { } } instructions { - name: "param_1.3" + name: "param_1.1" opcode: "parameter" shape { element_type: U32 @@ -508,9 +554,10 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "param_1.1" } parameter_number: 1 - id: 12884901891 + id: 8589934595 frontend_attributes { } statistics_viz { @@ -529,11 +576,12 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "data.2" } - id: 12884901888 - operand_ids: 12884901889 - operand_ids: 12884901890 - operand_ids: 12884901891 + id: 8589934592 + operand_ids: 8589934593 + operand_ids: 8589934594 + operand_ids: 8589934595 frontend_attributes { } statistics_viz { @@ -573,85 +621,14 @@ gpu_executable { } is_dynamic_dimension: false } - parameter_names: "param_0.4" - parameter_names: "param_1.3" + parameter_names: "param_0.1" + parameter_names: "param_1.1" parameter_names: "param_2.2" } - id: 3 - root_id: 12884901888 + id: 2 + root_id: 8589934592 is_fusion_computation: true } - computations { - name: "async_computation" - instructions { - name: "param_0.5" - opcode: "parameter" - shape { - element_type: U32 - dimensions: 8 - layout { - minor_to_major: 0 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - } - metadata { - } - id: 21474836480 - frontend_attributes { - } - statistics_viz { - } - } - instructions { - name: "ars.2" - opcode: "reduce-scatter" - shape { - element_type: U32 - dimensions: 4 - layout { - minor_to_major: 0 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - } - metadata { - } - dimensions: 0 - id: 21474836481 - operand_ids: 21474836480 - called_computation_ids: 0 - frontend_attributes { - } - statistics_viz { - } - collective_device_list { - } - } - program_shape { - parameters { - element_type: U32 - dimensions: 8 - layout { - minor_to_major: 0 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - } - result { - element_type: U32 - dimensions: 4 - layout { - minor_to_major: 0 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - } - parameter_names: "param_0.5" - } - id: 5 - root_id: 21474836481 - } computations { name: "main" instructions { @@ -742,6 +719,7 @@ gpu_executable { } } metadata { + scheduling_name: "id.1" } id: 4294967296 frontend_attributes { @@ -762,64 +740,22 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "loop_select_fusion" } fusion_kind: "kLoop" id: 4294967299 operand_ids: 4294967297 operand_ids: 4294967298 operand_ids: 4294967296 - called_computation_ids: 3 + called_computation_ids: 2 frontend_attributes { } statistics_viz { } - backend_config_payload { - id: 0 - } } instructions { - name: "reduce-scatter-start" - opcode: "async-start" - shape { - element_type: TUPLE - tuple_shapes { - element_type: TUPLE - tuple_shapes { - element_type: U32 - dimensions: 8 - layout { - minor_to_major: 0 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - } - } - tuple_shapes { - element_type: U32 - dimensions: 4 - layout { - minor_to_major: 0 - tail_padding_alignment_in_elements: 1 - } - is_dynamic_dimension: false - } - } - metadata { - } - id: 4294967300 - operand_ids: 4294967299 - called_computation_ids: 5 - frontend_attributes { - } - statistics_viz { - } - backend_config_payload { - id: 1 - } - } - instructions { - name: "reduce-scatter-done" - opcode: "async-done" + name: "ars.3" + opcode: "reduce-scatter" shape { element_type: U32 dimensions: 4 @@ -830,13 +766,21 @@ gpu_executable { is_dynamic_dimension: false } metadata { + scheduling_name: "reduce-scatter-start" } - id: 4294967301 - operand_ids: 4294967300 + dimensions: 0 + id: 4294967300 + operand_ids: 4294967299 + called_computation_ids: 0 frontend_attributes { } statistics_viz { } + collective_device_list { + } + backend_config_payload { + id: 0 + } } program_shape { result { @@ -850,7 +794,7 @@ gpu_executable { } } id: 1 - root_id: 4294967301 + root_id: 4294967300 } host_program_shape { result { @@ -881,14 +825,6 @@ gpu_executable { instruction_ids: 4294967297 instruction_ids: 4294967299 instruction_ids: 4294967300 - instruction_ids: 4294967301 - } - } - sequences { - key: 5 - value { - instruction_ids: 21474836480 - instruction_ids: 21474836481 } } } @@ -907,11 +843,10 @@ gpu_executable { frontend_attributes { map { key: "fingerprint_before_lhs" - value: "ecc7fa57c3eeb812403576a4054128b5" + value: "c36d8870be58d238935bebefbbc5b827" } } - payloads: "{\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"native_emitter_backend_config\":{\"type\":\"NATIVE_EMITTER_TYPE_INVALID\",\"unroll_factor\":0},\"operation_queue_id\":\"0\",\"reification_cost\":[]}" - payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" + payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}" } config { entry_computation_layout { @@ -942,7 +877,7 @@ gpu_executable { xla_cpu_enable_fast_math: false xla_gpu_enable_fast_min_max: false xla_force_host_platform_device_count: 1 - xla_gpu_enable_dynamic_slice_fusion: false + xla_gpu_enable_dynamic_slice_fusion: true xla_hlo_evaluator_use_fast_path: true xla_dump_hlo_as_html: false xla_cpu_fast_math_honor_nans: true @@ -968,7 +903,6 @@ gpu_executable { xla_gpu_use_runtime_fusion: false xla_dump_latency_hiding_schedule: false xla_dump_enable_mlir_pretty_form: true - xla_gpu_enable_latency_hiding_scheduler: false xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP xla_gpu_enable_triton_gemm: true xla_gpu_enable_cudnn_int8x32_convolution_reordering: true @@ -981,14 +915,10 @@ gpu_executable { xla_gpu_all_gather_combine_threshold_bytes: 31457287 xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287 xla_gpu_enable_highest_priority_async_stream: true - xla_gpu_enable_pipelined_all_reduce: false xla_gpu_exhaustive_tiling_search: false xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0 xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1 - xla_gpu_enable_pipelined_all_gather: false xla_gpu_redzone_padding_bytes: 8388608 - xla_gpu_enable_pipelined_reduce_scatter: true - xla_gpu_fused_attention_use_cudnn_rng: false xla_gpu_copy_insertion_use_region_analysis: false xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807 xla_gpu_collect_cost_model_stats: false @@ -1032,6 +962,7 @@ gpu_executable { xla_reduce_window_rewrite_base_length: 16 xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL xla_gpu_enable_host_memory_offloading: false + xla_gpu_exclude_nondeterministic_ops: true xla_gpu_nccl_terminate_on_error: false xla_gpu_shard_autotuning: true xla_gpu_enable_approx_costly_collectives: false @@ -1102,7 +1033,6 @@ gpu_executable { xla_cpu_generate_unique_c_style_kernel_entry_points: false xla_gpu_experimental_collective_cse_distance_threshold: 0 xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true - xla_cpu_use_fusion_emitters: true xla_gpu_experimental_collective_perf_table_path: "" xla_gpu_collective_permute_combine_threshold_bytes: 31457287 xla_unsupported_crash_on_hlo_pass_noop_change: false @@ -1111,7 +1041,6 @@ gpu_executable { xla_gpu_unsupported_enable_triton_multi_output_fusion: false xla_gpu_experimental_matmul_perf_table_path: "" xla_gpu_unsupported_enable_all_reduce_decomposer: false - xla_gpu_unsupported_use_all_reduce_one_shot_kernel: true xla_gpu_nccl_blocking_communicators: true xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20 xla_gpu_first_collective_call_terminate_timeout_seconds: 40 @@ -1122,14 +1051,13 @@ gpu_executable { xla_detect_unstable_reductions: DETECTION_MODE_NONE xla_gpu_command_buffer_scheduling_mode: LHS xla_gpu_experimental_enable_nccl_symmetric_buffers: false - xla_gpu_experimental_autotuner_cache_dir: "" xla_disable_automatic_host_compute_offload: false xla_gpu_experimental_enable_fusion_autotuner: true xla_gpu_experimental_scaled_dot_with_triton: true xla_gpu_command_buffer_unroll_loops: false xla_cpu_experimental_onednn_custom_call: false - xla_gpu_experimental_use_raft_select_k: false xla_gpu_experimental_enable_checksum_tracing_on_thunks: false + xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true xla_cpu_collective_call_terminate_timeout_seconds: 40 xla_cpu_collective_call_warn_stuck_seconds: 20 xla_keep_shardings_after_spmd: false @@ -1139,31 +1067,44 @@ gpu_executable { xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE xla_cpu_enable_platform_dependent_math: true xla_gpu_detect_nan: DETECTION_MODE_NONE - xla_gpu_experimental_dump_gpu_executable: true xla_gpu_detect_inf: DETECTION_MODE_NONE xla_enable_enzyme_comms_opt: false xla_gpu_experimental_enable_buffer_saver_on_thunks: false xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE - xla_gpu_experimental_aot_compiled_thunks: true xla_enable_scoped_logging_timers: true xla_gpu_executable_embed_debug_info: true xla_cpu_collective_timeout_seconds: 1800 xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false xla_gpu_enable_pipelined_host_offloading: false + xla_gpu_default_to_alg_dot_bf16_bf16_f32: false + xla_gpu_experimental_autotune_backends: CUDNN + xla_gpu_experimental_autotune_backends: TRITON + xla_gpu_experimental_autotune_backends: CUBLASLT + xla_gpu_experimental_autotune_backends: HIPBLASLT + xla_gpu_experimental_autotune_backends: MIOPEN + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER + xla_gpu_experimental_autotune_backends: NATIVE_EMITTER + xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER + xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION + xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION + xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION + xla_gpu_experimental_autotune_backends: YNNPACK + xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU xla_gpu_experimental_all_fusions_with_triton: false xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT xla_gpu_rocm_max_trace_events: 4194304 - xla_gpu_experimental_ragged_all_to_all_use_barrier: true xla_gpu_execution_terminate_timeout: "inf" - xla_enable_hlo_sharding_v3: false + xla_enable_hlo_sharding_v3: true xla_gpu_log_minmax: false xla_gpu_print_compilation_stats: false xla_gpu_experimental_enable_conv_fusion: false - xla_gpu_experimental_enable_tiling_propagation: false + xla_gpu_experimental_enable_tiling_propagation: true xla_gpu_execution_progress_tracking: 0 xla_gpu_experimental_max_unroll_factor: 32 xla_gpu_enable_pdl: true xla_gpu_collective_combine_threshold_count: 2048 + xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true xla_recognize_reduction_optimization_level: 0 xla_gpu_experimental_parallel_async_compute_limit: 2 xla_dump_buffer_assignment_analysis: true @@ -1183,12 +1124,12 @@ gpu_executable { xla_gpu_executable_num_compute_streams: 0 xla_gpu_executable_num_communication_streams: 0 xla_gpu_experimental_collective_start_as_early_as_possible: false - xla_cpu_experimental_enable_tiling_propagation: false + xla_cpu_experimental_enable_tiling_propagation: true xla_gpu_enable_dus_accumulator_zero_init_elimination: false xla_gpu_experimental_thunk_buffer_debug_module_outputs: false xla_cpu_ftz: true xla_gpu_enable_allocator_spatial_partitioning: true - xla_dump_compact_gte: false + xla_dump_compact_gte: true xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES xla_dump_hlo_as_riegeli: false xla_gpu_experimental_use_ragged_dot_grouped_gemm: true @@ -1197,6 +1138,39 @@ gpu_executable { xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED xla_gpu_enable_nccl_user_buffers_in_default_space: false xla_sdy_export_all_reduce_scatter: false + xla_gpu_enable_gxl_ragged_all_to_all: false + xla_gpu_gxl_scratch_size_bytes: 67108864 + xla_gpu_use_new_autotune_cache_format: true + xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE + xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT + xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON + xla_spmd_enable_dynamic_slice_collective_broadcast: false + xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false + xla_gpu_async_copy_min_bytes: -1 + xla_gpu_blas_max_algorithms: 0 + xla_gpu_experimental_enable_selective_memcpy_overlap: false + xla_gpu_experimental_enable_same_shape_multi_output_fusion: false + xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1 + xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1 + xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240 + xla_cpu_use_new_xtile_lowering: false + xla_gpu_trace_annotation_level: 0 + xla_gpu_experimental_cost_model_gemm_tiling_default: false + xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false + xla_gpu_all_reduce_splitter_ignore_profitability_check: false + xla_gpu_enable_cupti_multi_subscriber: true + xla_enable_hlo_modules_upload: true + xla_gpu_collective_domain_assignment: "" + xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1 + xla_compile_all_supported_configs: false + xla_deduplicate_backend_configs_min_size: 9223372036854775807 + xla_force_config: "" + xla_candidate_configs_file: "" + xla_gpu_enable_persistent_symmetric_memory: false + xla_gpu_experimental_enable_raft_for_stable_topk: false + xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND + xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET } static_device_assignment { replica_count: 2 @@ -1238,58 +1212,43 @@ gpu_executable { } logical_buffers { id: 10 - size: 16 - defined_at { - instruction_id: 21474836481 - } - } - logical_buffers { - id: 11 size: 32 defined_at { instruction_id: 4294967297 } } logical_buffers { - id: 12 + id: 11 size: 32 defined_at { instruction_id: 4294967298 } } logical_buffers { - id: 13 + id: 12 size: 4 defined_at { instruction_id: 4294967296 } } logical_buffers { - id: 14 + id: 13 size: 32 defined_at { instruction_id: 4294967299 } } logical_buffers { - id: 15 + id: 14 size: 16 defined_at { instruction_id: 4294967300 } } - logical_buffers { - id: 16 - size: 8 - defined_at { - shape_index: 0 - instruction_id: 4294967300 - } - } buffer_allocations { size: 32 assigned { - logical_buffer_id: 11 + logical_buffer_id: 10 size: 32 element_type: U32 } @@ -1299,7 +1258,7 @@ gpu_executable { index: 1 size: 32 assigned { - logical_buffer_id: 12 + logical_buffer_id: 11 size: 32 element_type: U32 } @@ -1310,13 +1269,13 @@ gpu_executable { size: 16 maybe_live_out: true assigned { - logical_buffer_id: 10 - size: 16 + logical_buffer_id: 12 + size: 4 element_type: U32 } assigned { - logical_buffer_id: 13 - size: 4 + logical_buffer_id: 14 + size: 16 element_type: U32 } } @@ -1335,7 +1294,6 @@ gpu_executable { size: 4 is_thread_local: true assigned { - logical_buffer_id: 1 size: 4 element_type: U32 } @@ -1345,56 +1303,45 @@ gpu_executable { size: 4 is_thread_local: true assigned { + logical_buffer_id: 1 size: 4 element_type: U32 } } buffer_allocations { index: 6 - size: 528 + size: 32 assigned { - logical_buffer_id: 14 + logical_buffer_id: 13 size: 32 element_type: U32 } - assigned { - logical_buffer_id: 15 - offset: 512 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 16 - offset: 256 - size: 8 - element_type: TUPLE - } } heap_simulator_traces { events { - buffer_id: 14 + buffer_id: 13 computation_name: "main" instruction_name: "loop_select_fusion" } events { kind: FREE - buffer_id: 14 + buffer_id: 13 computation_name: "main" instruction_name: "loop_select_fusion" } whole_module_simulation: true - buffer_allocation_index: 5 + buffer_allocation_index: 3 } peak_buffers { - logical_buffer_ids: 11 + logical_buffer_ids: 10 } peak_buffers { index: 1 - logical_buffer_ids: 12 + logical_buffer_ids: 11 } peak_buffers { index: 2 - logical_buffer_ids: 10 + logical_buffer_ids: 14 } peak_buffers { index: 3 @@ -1402,24 +1349,22 @@ gpu_executable { } peak_buffers { index: 4 - logical_buffer_ids: 1 + logical_buffer_ids: 0 } peak_buffers { index: 5 - logical_buffer_ids: 0 + logical_buffer_ids: 1 } peak_buffers { index: 6 - logical_buffer_ids: 14 - logical_buffer_ids: 16 - logical_buffer_ids: 15 + logical_buffer_ids: 13 } } - binary: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000@\005\000\000\000\000\000\000\000\003\000\000\000\000\000\000<\005<\000@\0008\000\004\000@\000\t\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.nv.global\000.nv.constant4\000.rel.nv.constant4\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.nv.global\000.nv.constant4\000.rel.nv.constant4\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000buffer_for_c1_1\000buffer_for_c0_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000=\000\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000y\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\210\000\000\000\021\000\010\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\230\000\000\000\021\000\010\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377s\000\000\000\000\000\000\000\000\000\000\021%\000\0056\010\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\210\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\250\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000]\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\002\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000y\000\000\000\013\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\002\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000K\000\000\000\t\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\002\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\007\000\000\000\010\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000=\000\000\000\001\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000h\002\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\003\000\000\000\000\000\000@\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\005\000\000\000@\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000h\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000@\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + binary: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000\000\004\000\000\000\000\000\000Z\rZ\000@\0008\000\004\000@\000\n\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.nv.shared.reserved.0\000.nv.global\000.nv.constant4\000\000\000\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant4\000\000\000\000.nv.callgraph\000.nv.prototype\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.global\000.nv.constant4\000\000\000\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant.pic\000.nv.callgraph\000.nv.prototype\000buffer_for_c1_1\000buffer_for_c0_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000a\000\000\000 \240\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\215\000\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\326\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\362\000\000\000\021\000\t\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\002\001\000\000\021\000\t\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\010\000\000\000\000\000\000\000\002\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\270\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\001\000\000\000\000\000\000\022\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\002\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000d\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\330\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\234\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\330\002\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\206\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\370\002\000\000\000\000\000\0000\000\000\000\000\000\000\000\003\000\000\000\007\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000S\000\000\000\001\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000(\003\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000@\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000(\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\0008\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\000\000\000\000\000\000\000\010\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000" gpu_compute_capability { cuda_compute_capability { - major: 6 - feature_extension: NONE + major: 9 + feature_extension: ACCELERATED_FEATURES } } module_name: "test" @@ -1437,87 +1382,36 @@ gpu_executable { buffer_allocations { values { size: 32 - assigned { - logical_buffer_id: 11 - size: 32 - element_type: U32 - } is_constant: true } values { index: 1 size: 32 - assigned { - logical_buffer_id: 12 - size: 32 - element_type: U32 - } is_constant: true } values { index: 2 size: 16 maybe_live_out: true - assigned { - logical_buffer_id: 10 - size: 16 - element_type: U32 - } - assigned { - logical_buffer_id: 13 - size: 4 - element_type: U32 - } } values { index: 3 size: 4 is_thread_local: true - assigned { - logical_buffer_id: 2 - size: 4 - element_type: U32 - } } values { index: 4 size: 4 is_thread_local: true - assigned { - logical_buffer_id: 1 - size: 4 - element_type: U32 - } } values { index: 5 size: 4 is_thread_local: true - assigned { - size: 4 - element_type: U32 - } } values { index: 6 - size: 528 - assigned { - logical_buffer_id: 14 - size: 32 - element_type: U32 - } - assigned { - logical_buffer_id: 15 - offset: 512 - size: 16 - element_type: TUPLE - } - assigned { - logical_buffer_id: 16 - offset: 256 - size: 8 - element_type: TUPLE - } + size: 32 } } output_info_map { @@ -1621,7 +1515,7 @@ gpu_executable { name: "loop_select_fusion" kernel_spec { cubin { - data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\000\006\000\000\000\000\000\000<\005<\000@\0008\000\003\000@\000\n\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.constant0.loop_select_fusion\000.rel.nv.constant0.loop_select_fusion\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.rel.nv.constant0.loop_select_fusion\000.nv.constant0.loop_select_fusion\000.nv.callgraph\000.nv.prototype\000.nv.rel.action\000loop_select_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\252\000\000\000\003\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\313\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\347\000\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\000\000\000\022\020\t\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\004/\010\000\005\000\000\000\007\000\000\000\004\021\010\000\005\000\000\000\000\000\000\000\004\022\010\000\005\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\0010\000\000\001*\000\000\004\n\010\000\002\000\000\000@\001 \000\003\031 \000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003\033\377\000\004\034\004\000\270\000\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377s\000\000\000\000\000\000\000\000\000\000\021%\000\0056\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\007\000\376\000D\034\000\001\000\207\000\200\007\230L\002\000G\005\200\007\230L\000\000\027\002\000\000\310\360\342\007 \366\000\204\037\000\003\000W\005\200\007\230L\002\002\007\000\000\240\324\356\004\000\007\005\200\007\230L\377\007 \376A\304\037\000\005\000\027\005\200\007\230L\006\000\'\000\000\000H8\000\000\347\001\000\000(8\361\007\251\374\000\230\037\000\004\004\'\005\000\001DK\005\0057\005\000\001DK\004\004g\000\000\200\020\\\342\007 \366\000\230\037\000\005\005\007\000\000\010\020\\\004\004\007\000\000\240\324\356\002\006g\005\000\200\020L\347\007 \376 \374\037\000\003\000w\005\000\010\020L\004\002\007\000\000 \334\356\017\000\007\000\000\000\000\343\377\007\000\374\000\200\037\000\017\000\007\377\377\017@\342\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\340\007\000\374\000\200\037\000\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\017\007\000\000\000\260P\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\366\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0006\001\000\000\000\000\000\000\t\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\002\000\000\000\000\000\000\220\000\000\000\000\000\000\000\002\000\000\000\005\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\002\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000K\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\364\002\000\000\000\000\000\000|\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\313\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\003\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\347\000\000\000\013\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\205\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\003\000\000\000\000\000\000`\001\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\000\001\000\000\000\000\000\000\003\000\000\000\005\000\000\007 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\005\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\240\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000`\002\000\000\000\000\000\000`\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\250\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000" + data: "\177ELF\002\001\0013\007\000\000\000\000\000\000\000\002\000\276\000\200\000\000\000\000\000\000\000\000\000\000\0000\014\000\000\000\000\000\0000\t\000\000\000\000\000\000Z\rZ\000@\0008\000\005\000@\000\014\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_select_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.nv.info\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_select_fusion\000.nv.constant0.loop_select_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0002\000\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\205\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\264\000\000\000 \240\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\325\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\001\000\000\003\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000!\001\000\000\022\020\t\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\0004\001\000\000\003\000\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\004<\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\004/\010\000\006\000\000\000\014\000\000\000\004\021\010\000\006\000\000\000\000\000\000\000\004\022\010\000\006\000\000\000\000\000\000\000\0047\004\000\200\000\000\000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\004\034\004\000\360\000\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000\003\031 \000\004\n\010\000\007\000\000\000\020\002 \000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\000\000\000\000D\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\n\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\210\000\000\000\n\000\000\000\"\016\000\271z\004\000\000\202\000\000\000\n\000\000\000\316\017\000\202{\004\377\000\204\000\000\000\n\000\000\000b\016\000\031y\t\000\000\000\000\000\000!\000\000\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\314\017\000\202{\006\377\000\212\000\000\000\n\000\000\000\342\016\000\201y\002\002\004\000\000\000\000\231\036\014\000$\037\000\014r\000\002\377\000\000\000pP\360\003\000\332\017\001\202\013\004\377\000\206\000\000\000\010\000\000\000\260.\000\202\013\005\377\000\207\000\000\000\010\000\000\000\244\016\000%x\004\t\004\000\000\000\004\000\216\007\000\314O\000\201y\005\004\004\000\000\000\000\231\036\014\000\242\016\000%x\006\t\004\000\000\000\006\000\216\007\000\312\217\000\206y\000\006\005\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000z\001\000\000\000\000\000\000U\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\002\000\000\000\000\000\000\300\000\000\000\000\000\000\000\002\000\000\000\006\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\233\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\370\003\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000K\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\034\004\000\000\000\000\000\000\200\000\000\000\000\000\000\000\003\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\313\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\234\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\271\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\004\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\0002\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000\000\002\000\000\000\000\000\000\003\000\000\000\006\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\205\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\347\000\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\0000\002\000\000\000\000\000\000\000\000\000\000\t\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\0000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\0000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\000\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0000\002\000\000\000\000\000\0000\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000" } arity: 4 kernel_name: "loop_select_fusion" @@ -1666,13 +1560,14 @@ gpu_executable { } } } + use_pdl: true tma_metadata { } } } thunks { thunk_info { - profile_annotation: "ars.2" + profile_annotation: "ars.3" thunk_id: 3 } reduce_scatter_thunk { @@ -1721,9 +1616,9 @@ gpu_executable { platform_name: "CUDA" cuda_platform_version { cuda_toolkit_version: "12.8.0" - cudnn_version: "9.10.0" + cudnn_version: "9.10.2" cub_version: "3.1.2" } } - buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 64B\n maybe_live_out allocation: 16B\n preallocated temp allocation: 528B\n preallocated temp fragmentation: 0B (0.00%)\n total allocation: 620B\n total fragmentation: 548B (88.39%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 32B\n\t\tXLA Label: fusion\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: async-start\n\t\tShape: ((u32[8]), u32[4])\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 16B\n\t\tXLA Label: reduce-scatter\n\t\tShape: u32[4]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 8B\n\t\tXLA Label: async-start\n\t\tShape: (u32[8])\n\t\t==========================\n\n\tBuffer 7:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 8:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 9:\n\t\tSize: 4B\n\t\tXLA Label: add\n\t\tShape: u32[]\n\t\t==========================\n" + buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n parameter allocation: 0B\n constant allocation: 64B\n maybe_live_out allocation: 16B\n preallocated temp allocation: 32B\n preallocated temp fragmentation: 0B (0.00%)\n total allocation: 124B\n total fragmentation: 76B (61.29%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 32B\n\t\tXLA Label: fusion\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: reduce-scatter\n\t\tShape: u32[4]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 7:\n\t\tSize: 4B\n\t\tXLA Label: add\n\t\tShape: u32[]\n\t\t==========================\n" } diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.cc b/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.cc index 98d1c661a30001..99c74ff514a656 100644 --- a/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.cc +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.cc @@ -23,13 +23,19 @@ limitations under the License. #include #include +#include "absl/log/check.h" +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/statusor.h" +#include "absl/strings/ascii.h" #include "absl/strings/match.h" #include "absl/strings/numbers.h" #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "xla/pjrt/pjrt_client.h" +#include "xla/stream_executor/platform.h" +#include "xla/stream_executor/platform_manager.h" +#include "xla/stream_executor/stream_executor.h" #include "xla/tests/aot_interception_pjrt_client.h" #include "xla/tests/hlo_test_base.h" #include "xla/tests/pjrt_client_registry.h" @@ -71,20 +77,53 @@ bool IsEnvFlagEnabled(const char* name) { } // The undeclared-outputs directory exists only when the golden-update helper +// (`third_party/tensorflow/compiler/xla/tests/aot_compatibility_experimental/google/update_goldens.py`) // drives the test. Shared so both callers enforce it with the same message. absl::StatusOr GetUndeclaredOutputsDir() { const char* out_dir = std::getenv("TEST_UNDECLARED_OUTPUTS_DIR"); if (out_dir == nullptr) { return absl::FailedPreconditionError( "TEST_UNDECLARED_OUTPUTS_DIR is unset. Golden updates must be driven " - "by the golden-update helper for this package, which sets it; running " - "the test target directly will not work."); + "by the golden-update helper for this package " + "(Google-internal: third_party/tensorflow/compiler/xla/tests/" + "aot_compatibility_experimental/google/update_goldens.py), which sets " + "it; running the test target directly will not work."); } return std::string(out_dir); } } // namespace +std::string DetectGpuArchToken() { + if (const char* env_arch = std::getenv("XLA_AOT_GOLDEN_ARCH")) { + return env_arch; + } + absl::StatusOr platform = + stream_executor::PlatformManager::PlatformWithName("CUDA"); + CHECK_OK(platform.status()) + << "AOT golden arch detection failed to get the CUDA platform; set " + "XLA_AOT_GOLDEN_ARCH to override."; + absl::StatusOr executor = + (*platform)->ExecutorForDevice(0); + CHECK_OK(executor.status()) + << "AOT golden arch detection failed to get a device executor; set " + "XLA_AOT_GOLDEN_ARCH to override."; + const std::string name = + absl::AsciiStrToLower((*executor)->GetDeviceDescription().name()); + + // Longest token first: "gb200" contains "b200". gb300 is deliberately + // absent -- it is listed in gpu/BUILD's disabled_backends (b/491194726), so + // no gb300 goldens exist and no gb300 target is generated. + for (const absl::string_view arch : {"gb200", "b200", "h100"}) { + if (absl::StrContains(name, arch)) { + return std::string(arch); + } + } + + LOG(FATAL) << "Unrecognized GPU arch for AOT goldens: " << name + << "; set XLA_AOT_GOLDEN_ARCH to override."; +} + std::string GetExecutablesDirectory(absl::string_view target_name, AOTTestPlatform platform) { return tsl::io::JoinPath( diff --git a/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.h b/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.h index 5f0b89b9bd71bb..43a6fe9dad2a46 100644 --- a/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.h +++ b/third_party/xla/xla/tests/aot_compatibility_experimental/test_lib.h @@ -48,6 +48,13 @@ struct AotTestParam { } }; +// Returns the golden-directory token for the GPU this process is running on: +// one of "h100", "b200", "gb200". `XLA_AOT_GOLDEN_ARCH`, when set, is returned +// verbatim and skips device detection; this is what update_goldens.py uses. +// Crashes with LOG(FATAL) if the device cannot be queried or its name matches +// no known token. +std::string DetectGpuArchToken(); + // Returns the path to the executables directory for the given test target and // platform. std::string GetExecutablesDirectory(absl::string_view target_name, diff --git a/third_party/xla/xla/tests/aot_interception_pjrt_client.cc b/third_party/xla/xla/tests/aot_interception_pjrt_client.cc index b308f43752c725..5b1fab26b8c1da 100644 --- a/third_party/xla/xla/tests/aot_interception_pjrt_client.cc +++ b/third_party/xla/xla/tests/aot_interception_pjrt_client.cc @@ -16,6 +16,7 @@ limitations under the License. #include "xla/tests/aot_interception_pjrt_client.h" #include +#include #include #include #include @@ -119,6 +120,23 @@ absl::Status CompareStructurally( return absl::OkStatus(); } +std::string RegenerationHint(absl::string_view target_name) { + // Bazel exports the running test's full label as TEST_TARGET, which is what + // update_goldens.py takes. Fall back to a placeholder label if it is unset. + const char* test_target = std::getenv("TEST_TARGET"); + std::string label = test_target != nullptr + ? std::string(test_target) + : absl::StrCat("//:", target_name); + return absl::StrCat( + "\n\nRegenerate the goldens for target `", label, + "` and add the newly generated v directory to the change after " + "reviewing the differences.\n" + "(Google-internal: run, from the workspace root, " + "third_party/tensorflow/compiler/xla/tests/" + "aot_compatibility_experimental/google/update_goldens.py ", + label, ")"); +} + } // namespace absl::StatusOr @@ -271,12 +289,9 @@ absl::Status AOTInterceptionPjrtClient::VerifyAgainstGolden( tsl::io::Basename(tsl::io::Dirname(tsl::io::Dirname(artifact_path_))); return absl::Status( diff_status.code(), - absl::StrCat( - "Golden artifact verification failed for ", artifact_path_, ".\n", - diff_status.message(), "\n\nRegenerate the goldens for target `", - target_name, - "` and add the newly generated v directory to the change " - "after reviewing the differences.")); + absl::StrCat("Golden artifact verification failed for ", artifact_path_, + ".\n", diff_status.message(), + RegenerationHint(target_name))); } return absl::OkStatus(); diff --git a/third_party/xla/xla/tools/tests/hlo_expand_test.cc b/third_party/xla/xla/tools/tests/hlo_expand_test.cc index a38dd367957162..7365f107b8ba94 100644 --- a/third_party/xla/xla/tools/tests/hlo_expand_test.cc +++ b/third_party/xla/xla/tools/tests/hlo_expand_test.cc @@ -88,9 +88,9 @@ TEST_F(HloExpandTest, SpmdHloWithoutHloShardingV3) { R"(HloModule module, entry_computation_layout={(f32[24,64]{1,0}, f32[39296,64]{1,0})->f32[24,19648]{1,0}}, num_partitions=2 ENTRY %entry_spmd (param: f32[24,64], param.1: f32[39296,64]) -> f32[24,19648] { - %param = f32[24,64]{1,0} parameter(0), sharding={replicated} + %param = f32[24,64]{1,0} parameter(0) %lhs.copy.1 = f32[24,64]{1,0} copy(%param) - %param.1 = f32[39296,64]{1,0} parameter(1), sharding={replicated} + %param.1 = f32[39296,64]{1,0} parameter(1) %constant = s32[2]{0} constant({0, 19648}) %partition-id = u32[] partition-id() %dynamic-slice = s32[1]{0} dynamic-slice(%constant, %partition-id), dynamic_slice_sizes={1} @@ -118,9 +118,9 @@ TEST_F(HloExpandTest, SpmdHloWithHloShardingV3) { R"(HloModule module, entry_computation_layout={(f32[24,64]{1,0}, f32[39296,64]{1,0})->f32[24,19648]{1,0}}, num_partitions=2 ENTRY %entry_spmd (param: f32[24,64], param.1: f32[39296,64]) -> f32[24,19648] { - %param = f32[24,64]{1,0} parameter(0), sharding={mesh[], replicated} + %param = f32[24,64]{1,0} parameter(0) %lhs.copy.1 = f32[24,64]{1,0} copy(%param) - %param.1 = f32[39296,64]{1,0} parameter(1), sharding={mesh[], replicated} + %param.1 = f32[39296,64]{1,0} parameter(1) %constant = s32[2]{0} constant({0, 19648}) %partition-id = u32[] partition-id() %dynamic-slice = s32[1]{0} dynamic-slice(%constant, %partition-id), dynamic_slice_sizes={1} diff --git a/third_party/xla/xla/tsl/concurrency/future.h b/third_party/xla/xla/tsl/concurrency/future.h index aca4ac089c050e..487057effb65f5 100644 --- a/third_party/xla/xla/tsl/concurrency/future.h +++ b/third_party/xla/xla/tsl/concurrency/future.h @@ -1719,7 +1719,7 @@ Future> JoinFutures(absl::Span> futures) { VLOG(2) << "tsl::JoinFutures: " << futures.size() << " futures"; if (futures.empty()) { - return Future>({}); + return Future>(std::vector{}); } auto [promise, future] = MakePromise>(); @@ -1743,7 +1743,7 @@ Future> JoinFutures(absl::Span> futures) { VLOG(2) << "tsl::JoinFutures: " << futures.size() << " futures"; if (futures.empty()) { - return Future>({}); + return Future>(std::vector{}); } auto [promise, future] = MakePromise>(); diff --git a/third_party/xla/xla/tsl/concurrency/future_test.cc b/third_party/xla/xla/tsl/concurrency/future_test.cc index 43a71c6e6b268d..d60575ffc83cd7 100644 --- a/third_party/xla/xla/tsl/concurrency/future_test.cc +++ b/third_party/xla/xla/tsl/concurrency/future_test.cc @@ -1087,6 +1087,25 @@ TEST(FutureTest, JoinCopyableFuturesError) { EXPECT_EQ(join_two.Await().status(), absl::InternalError("error0")); } +TEST(FutureTest, JoinEmptyCopyableFutures) { + std::vector> futures; + Future> join = JoinFutures(futures); + ASSERT_TRUE(join.IsValid()); + EXPECT_TRUE(join.IsReady()); + ASSERT_OK_AND_ASSIGN(std::vector v, join.Await()); + EXPECT_TRUE(v.empty()); +} + +TEST(FutureTest, JoinEmptyMoveOnlyFutures) { + std::vector>> futures; + Future>> join = + JoinFutures>(absl::MakeSpan(futures)); + ASSERT_TRUE(join.IsValid()); + EXPECT_TRUE(join.IsReady()); + ASSERT_OK_AND_ASSIGN(auto v, std::move(join).Await()); + EXPECT_TRUE(v.empty()); +} + TEST(FutureTest, JoinMoveOnlyFuture) { auto [promise0, future0] = MakePromise>(); auto [promise1, future1] = MakePromise>(); diff --git a/third_party/xla/xla/tsl/platform/default/env.cc b/third_party/xla/xla/tsl/platform/default/env.cc index 6ed6443235f56d..4b52100e7d6c06 100644 --- a/third_party/xla/xla/tsl/platform/default/env.cc +++ b/third_party/xla/xla/tsl/platform/default/env.cc @@ -206,7 +206,7 @@ class PosixEnv : public Env { } void SchedClosure(absl::AnyInvocable closure) override { - // TODO(b/27290852): Spawning a new thread here is wasteful, but + // TODO(b/119634930): Spawning a new thread here is wasteful, but // needed to deal with the fact that many `closure` functions are // blocking in the current codebase. std::thread closure_thread(std::move(closure)); @@ -215,7 +215,7 @@ class PosixEnv : public Env { void SchedClosureAfter(int64_t micros, absl::AnyInvocable closure) override { - // TODO(b/27290852): Consuming a thread here is wasteful, but this + // TODO(b/119634930): Consuming a thread here is wasteful, but this // code is (currently) only used in the case where a step fails // (AbortStep). This could be replaced by a timer thread SchedClosure([this, micros, closure = std::move(closure)]() mutable { diff --git a/third_party/xla/xla/tuple_tree.h b/third_party/xla/xla/tuple_tree.h index 75671f26320a3f..37e7668bb05fb5 100644 --- a/third_party/xla/xla/tuple_tree.h +++ b/third_party/xla/xla/tuple_tree.h @@ -357,6 +357,17 @@ class TupleTree { return entry_or.value()->children_start_id == -1; } + // Returns the number of children of the node at the given index, or 0 if the + // node does not exist or is a leaf. + size_t num_children(ShapeIndexView index = {}) const { + absl::StatusOr entry_or = + index_table_.GetEntry(index); + if (!entry_or.ok() || entry_or.value()->children_start_id == -1) { + return 0; + } + return entry_or.value()->num_children; + } + // Checks if the structure of this TupleTree is compatible with the given // shape. bool IsStructurallyCompatible(const Shape& shape) const { diff --git a/third_party/xla/xla/tuple_tree_test.cc b/third_party/xla/xla/tuple_tree_test.cc index 5741d694424682..353124d76123ae 100644 --- a/third_party/xla/xla/tuple_tree_test.cc +++ b/third_party/xla/xla/tuple_tree_test.cc @@ -46,12 +46,14 @@ namespace { TEST_F(TupleTreeTest, SingleLeafConstructor) { TupleTree tree(42); EXPECT_TRUE(tree.IsLeaf({})); + EXPECT_EQ(tree.num_children(), 0); EXPECT_EQ(tree.element({}), 42); } TEST_F(TupleTreeTest, EmptyConstructor) { TupleTree tree; EXPECT_FALSE(tree.IsLeaf({})); + EXPECT_EQ(tree.num_children(), 0); EXPECT_THAT(tree.nodes(), ElementsAre(Pair(ShapeIndex({}), 0))); EXPECT_THAT(tree.leaves(), ElementsAre()); EXPECT_EQ(tree.element({}), 0); @@ -71,12 +73,15 @@ TEST_F(TupleTreeTest, NodeWithValueAndChildren) { // Node with value and empty children (empty tuple) TupleTree tree1(Node::Tuple(100)); EXPECT_FALSE(tree1.IsLeaf({})); // It's a tuple, not a leaf + EXPECT_EQ(tree1.num_children(), 0); EXPECT_THAT(tree1.nodes(), ElementsAre(Pair(ShapeIndex({}), 100))); EXPECT_THAT(tree1.leaves(), ElementsAre()); // Node with value and non-empty children TupleTree tree2(Node::Tuple(200, {Node::Leaf(1), Node::Leaf(2)})); EXPECT_FALSE(tree2.IsLeaf({})); + EXPECT_EQ(tree2.num_children(), 2); + EXPECT_EQ(tree2.num_children({0}), 0); EXPECT_TRUE(tree2.IsLeaf({0})); EXPECT_EQ(tree2.element({0}), 1); EXPECT_TRUE(tree2.IsLeaf({1})); diff --git a/third_party/xla/xla/util/split_proto/BUILD b/third_party/xla/xla/util/split_proto/BUILD index 3f8d0dc579cff4..8a5122362547d3 100644 --- a/third_party/xla/xla/util/split_proto/BUILD +++ b/third_party/xla/xla/util/split_proto/BUILD @@ -74,7 +74,9 @@ cc_library( "@com_google_absl//absl/strings:string_view", "@com_google_protobuf//:protobuf", "@riegeli//riegeli/base:any", + "@riegeli//riegeli/base:chain", "@riegeli//riegeli/bytes:reader", + "@riegeli//riegeli/messages:parse_message", "@riegeli//riegeli/records:record_position", "@riegeli//riegeli/records:record_reader", ], diff --git a/third_party/xla/xla/util/split_proto/split_executable_and_options_writer.cc b/third_party/xla/xla/util/split_proto/split_executable_and_options_writer.cc index 8ef9c2ba1779c3..828373418e79ef 100644 --- a/third_party/xla/xla/util/split_proto/split_executable_and_options_writer.cc +++ b/third_party/xla/xla/util/split_proto/split_executable_and_options_writer.cc @@ -78,8 +78,13 @@ ExecutableAndOptionsProto GetProtoWithoutSerializedExecutable( absl::Status WriteSplitExecutableAndOptions( const ExecutableAndOptionsProto& executable_and_options, std::unique_ptr writer) { - riegeli::RecordWriter record_writer(std::move(writer), - GetGpuSplitProtoOptions()); + riegeli::RecordWriter record_writer( + std::move(writer), + // The bulk of this proto is `serialized_executable`, which is already a + // compressed split proto. Compressing it again does not reduce the size + // but adds a full (de)compression pass over the whole executable when + // serializing and when loading it. + GetSplitProtoRiegeliOptions(SplitProtoCompression::kNone)); SplitProtoManifest manifest = BuildManifest(); TF_RETURN_WITH_CONTEXT_IF_ERROR( WriteRecord(record_writer, manifest), diff --git a/third_party/xla/xla/util/split_proto/split_proto_reader.cc b/third_party/xla/xla/util/split_proto/split_proto_reader.cc index 47fc1419aa08ab..3e3a18284ccf3f 100644 --- a/third_party/xla/xla/util/split_proto/split_proto_reader.cc +++ b/third_party/xla/xla/util/split_proto/split_proto_reader.cc @@ -27,7 +27,9 @@ limitations under the License. #include "google/protobuf/message.h" #include "google/protobuf/reflection.h" #include "riegeli/base/any.h" +#include "riegeli/base/chain.h" #include "riegeli/bytes/reader.h" +#include "riegeli/messages/parse_message.h" #include "riegeli/records/record_position.h" #include "riegeli/records/record_reader.h" #include "xla/tsl/platform/errors.h" @@ -52,15 +54,20 @@ absl::Status ReadRecord(riegeli::RecordReader& record_reader, T& record) { template absl::Status HandleProtoMergeRecord(riegeli::RecordReader& record_reader, google::protobuf::Message& proto) { - absl::string_view record_data; + // Read the record as a `Chain` rather than a flat `string_view`: large + // records (e.g. HLO modules with big constants) span multiple blocks of the + // decoded chunk, and reading them as a flat view would force an extra copy + // to make them contiguous. The `Chain` shares the decoded blocks instead and + // the proto parser can consume it directly. + riegeli::Chain record_data; TF_RETURN_WITH_CONTEXT_IF_ERROR(ReadRecord(record_reader, record_data), "failed to read proto merge record data for ", proto.GetTypeName()); - if (!proto.MergeFromString(record_data)) { - return absl::InternalError(absl::StrFormat( - "Failed to parse proto merge record for %s", proto.GetTypeName())); - } + TF_RETURN_WITH_CONTEXT_IF_ERROR( + riegeli::ParseMessage(record_data, proto, + riegeli::ParseMessageOptions().set_merge(true)), + "Failed to parse proto merge record for ", proto.GetTypeName()); return absl::OkStatus(); } diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto index 36f544733c1151..44ef782f522cc1 100644 --- a/third_party/xla/xla/xla.proto +++ b/third_party/xla/xla/xla.proto @@ -1867,7 +1867,7 @@ message DebugOptions { // for all instructions. Accepts textproto of autotuner.CandidateConfigs. optional string xla_candidate_configs_file = 544; - // If true, allows skipping defensive copy insertion for S(1) collective + // If true, allows skipping defensive copy insertion for collective // memory parameters that have input-output aliasing and execute on all // available devices in the topology. optional bool xla_gpu_enable_persistent_symmetric_memory = 545; diff --git a/third_party/xla/xla/xla_data.proto b/third_party/xla/xla/xla_data.proto index a8465a05da565d..9cd793c6e2a362 100644 --- a/third_party/xla/xla/xla_data.proto +++ b/third_party/xla/xla/xla_data.proto @@ -965,6 +965,35 @@ enum RandomAlgorithm { // Next: 2 } +// The mode of a `shuffle` instruction together with the attributes that +// parameterize that mode. A shuffle moves the elements of its operand around +// along its shuffled dimensions, and the mode is the pattern of that movement. +// Every mode owns a message holding exactly the attributes that the mode +// needs, and the active case of the `mode` oneof selects the mode; a mode +// without attributes owns an empty message. +// +// The name of a case of the `mode` oneof is also the name of the mode in the +// HLO text format, e.g. `mode=rotate`. +message ShuffleMode { + // Rotates the elements along each shuffled dimension to the left, wrapping + // the elements that fall off the front around to the back, e.g. rotating + // `[1, 2, 3, 4]` by a shift of 1 yields `[2, 3, 4, 1]`. + message Rotate { + // The shift of each rotated dimension, in the order in which the + // dimensions appear in `dimensions`. A shift of `s` moves the + // element at index `i` of a dimension of size `d` to index + // `(i - s) mod d`, so `s` is the index of the element that ends up first. + // Shifts outside of `[0, d)`, including negative ones, wrap around. + repeated int64 shifts = 1; + } + + oneof mode { + Rotate rotate = 1; + } + + // Next: 2 +} + message TriangularSolveOptions { // If true, solves ax = b. If false, solves xa = b. bool left_side = 1;