From bbecf0368e0f5c386a153daa706cf150736aec91 Mon Sep 17 00:00:00 2001 From: prasanna8585 Date: Thu, 30 Jul 2026 20:36:49 +0530 Subject: [PATCH 01/32] Fix code injection via unvalidated Op attr/arg names (GHSA-2xp6-8g4h-qw72) ValidateOpDef()/ValidateArg() in op_def_util.cc validated the op's own name against IsValidOpName but never validated the character set of attribute names (OpDef.AttrDef.name()) or argument names (OpDef.ArgDef.name()) -- only their uniqueness. TensorFlow's op-wrapper code generators (cc_op_gen.cc, python_op_gen.cc) splice these names directly into generated C++/Python source with no escaping. A crafted attr/arg name containing ", ), ;, or a comma can break out of its syntactic context and inject arbitrary code, executed at compile time (C++) or import time (Python) of the generated wrapper. This is the same class of bug already fixed in the sibling file cc_op_fuzz_gen.cc (escaping via absl::CEscape). This change: 1. Adds IsValidAttrOrArgName(), a safe-identifier check (must start with a letter/underscore, contain only letters/digits/underscores), and applies it in ValidateOpDef() and ValidateArg(). This closes the injection class at the single validation choke-point, so it protects every current and future downstream consumer -- not just the two known generators. 2. Applies absl::CEscape() to graph_attr.name() at the cc_op_gen.cc splice site as defense-in-depth, matching the precedent already merged in cc_op_fuzz_gen.cc. Fixes GHSA-2xp6-8g4h-qw72. --- tensorflow/cc/framework/cc_op_gen.cc | 5 ++-- tensorflow/core/framework/op_def_util.cc | 30 ++++++++++++++++++++++++ 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/tensorflow/cc/framework/cc_op_gen.cc b/tensorflow/cc/framework/cc_op_gen.cc index 813cba37dcce43..fb4e752c9b9533 100644 --- a/tensorflow/cc/framework/cc_op_gen.cc +++ b/tensorflow/cc/framework/cc_op_gen.cc @@ -241,8 +241,9 @@ std::string GetConstructorBody(const OpInfo& op_info) { api_def_attr.has_default_value() ? absl::StrCat("attrs.", api_def_attr.rename_to(), "_") : AvoidCPPKeywords(api_def_attr.rename_to()); - strings::StrAppend(&body, spaces, ".Attr(\"", graph_attr.name(), "\", ", - attr_name, ")\n"); + strings::StrAppend(&body, spaces, ".Attr(\"", + absl::CEscape(graph_attr.name()), "\", ", attr_name, + ")\n"); } absl::StrAppend(&body, " ;\n"); absl::StrAppend(&body, " ", scope_str, ".UpdateBuilder(&builder);\n"); diff --git a/tensorflow/core/framework/op_def_util.cc b/tensorflow/core/framework/op_def_util.cc index 0a3018a3863bb3..dba84dbbfecc97 100644 --- a/tensorflow/core/framework/op_def_util.cc +++ b/tensorflow/core/framework/op_def_util.cc @@ -191,12 +191,19 @@ const ApiDef::Arg* FindInputArg(absl::string_view name, const ApiDef& api_def) { } \ } while (false) +// Forward declaration; defined below alongside IsValidOpName. +bool IsValidAttrOrArgName(absl::string_view sp); + static absl::Status ValidateArg(const OpDef::ArgDef& arg, const OpDef& op_def, bool output, absl::flat_hash_set* names) { const std::string suffix = absl::StrCat(output ? " for output '" : " for input '", arg.name(), "'"); VALIDATE(names->emplace(arg.name()).second, "Duplicate name: ", arg.name()); + VALIDATE(IsValidAttrOrArgName(arg.name()), "Invalid argument name: ", + arg.name(), + " (must start with a letter or underscore and contain only " + "letters, digits, and underscores)"); VALIDATE(HasAttrStyleType(arg), "Missing type", suffix); if (!arg.number_attr().empty()) { @@ -267,6 +274,25 @@ bool IsValidOpName(absl::string_view sp) { } } +// Attribute and argument names are not restricted to CamelCase like op +// names, but they are spliced verbatim (unescaped) into generated C++ and +// Python wrapper source by tools such as cc_op_gen.cc and python_op_gen.cc. +// To prevent a malicious attr/arg name from injecting arbitrary code into +// that generated source (e.g. a name containing `", ), ;`), restrict these +// names to a safe identifier character set: they must start with a letter +// or underscore and contain only letters, digits, and underscores. +bool IsValidAttrOrArgName(absl::string_view sp) { + using ::tensorflow::strings::Scanner; + + if (sp.empty()) return false; + if (sp[0] != '_' && !isalpha(static_cast(sp[0]))) { + return false; + } + Scanner scanner(sp); + scanner.Any(Scanner::LETTER_DIGIT_UNDERSCORE); + return scanner.GetResult() && scanner.empty(); +} + absl::Status ValidateOpDef(const OpDef& op_def) { if (!absl::StartsWith(op_def.name(), "_")) { VALIDATE(IsValidOpName(op_def.name()), "Invalid name: ", op_def.name(), @@ -279,6 +305,10 @@ absl::Status ValidateOpDef(const OpDef& op_def) { // Validate name VALIDATE(names.emplace(attr.name()).second, "Duplicate name: ", attr.name()); + VALIDATE(IsValidAttrOrArgName(attr.name()), "Invalid attr name: ", + attr.name(), + " (must start with a letter or underscore and contain only " + "letters, digits, and underscores)"); DataType dt; VALIDATE(!DataTypeFromString(attr.name(), &dt), "Attr can't have name ", attr.name(), " that matches a data type"); From 0ffb08a42cf3586e0be5fe58e12f69576abb0a9f Mon Sep 17 00:00:00 2001 From: prasanna8585 Date: Thu, 30 Jul 2026 20:56:55 +0530 Subject: [PATCH 02/32] Address review: use absl::ascii_isalpha/ascii_isalnum in IsValidAttrOrArgName Per gemini-code-assist review feedback on PR #124374: std::isalpha is locale-dependent and can behave unexpectedly depending on environment locale settings, and mixing it with Scanner's strictly-ASCII character classes was inconsistent. Replaces both with Abseil's locale-independent ASCII utility functions (absl::ascii_isalpha, absl::ascii_isalnum), which also simplifies the function and drops the Scanner dependency entirely. No behavioral change: verified against the same test vectors as before (malicious attr name from the advisory rejected; legitimate names like dtype, T, num_threads_2, _internal still accepted). --- tensorflow/core/framework/op_def_util.cc | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/tensorflow/core/framework/op_def_util.cc b/tensorflow/core/framework/op_def_util.cc index dba84dbbfecc97..64f82f89516cbe 100644 --- a/tensorflow/core/framework/op_def_util.cc +++ b/tensorflow/core/framework/op_def_util.cc @@ -21,6 +21,7 @@ limitations under the License. #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" +#include "absl/strings/ascii.h" #include "tensorflow/core/framework/attr_value.pb.h" #include "tensorflow/core/framework/attr_value_util.h" #include "tensorflow/core/framework/op_def.pb.h" @@ -282,15 +283,17 @@ bool IsValidOpName(absl::string_view sp) { // names to a safe identifier character set: they must start with a letter // or underscore and contain only letters, digits, and underscores. bool IsValidAttrOrArgName(absl::string_view sp) { - using ::tensorflow::strings::Scanner; - if (sp.empty()) return false; - if (sp[0] != '_' && !isalpha(static_cast(sp[0]))) { + if (sp[0] != '_' && !absl::ascii_isalpha(sp[0])) { return false; } - Scanner scanner(sp); - scanner.Any(Scanner::LETTER_DIGIT_UNDERSCORE); - return scanner.GetResult() && scanner.empty(); + for (size_t i = 1; i < sp.size(); ++i) { + char c = sp[i]; + if (c != '_' && !absl::ascii_isalnum(c)) { + return false; + } + } + return true; } absl::Status ValidateOpDef(const OpDef& op_def) { From 3ed42b0d486a31a37416fcea47f6212a753fcc43 Mon Sep 17 00:00:00 2001 From: prasanna8585 <65734642+prasanna8585@users.noreply.github.com> Date: Mon, 3 Aug 2026 09:29:31 +0530 Subject: [PATCH 03/32] Address review: validate ApiDef rename_to(), add regression tests Per review feedback on this PR: 1. The PR description claimed defense-in-depth escaping was applied to both generators, but python_op_gen.cc was untouched. Investigating why revealed the real issue: both generators splice names into IDENTIFIER positions (C++ parameter names, Python keyword argument names) via ApiDef.Arg/Attr.rename_to() -- a field entirely separate from the OpDef.ArgDef/AttrDef.name() fields validated by ValidateOpDef/ValidateArg. CEscape (for string literals) doesn't apply to identifier positions at all, so 'add escaping to python_op_gen.cc' was the wrong frame -- rename_to() needed validation, not escaping, in both generators. 2. This also directly addresses the GetConstructorDecl() case flagged in review: op_info.arg_names[i] traces to AvoidCPPKeywords(api_def_arg.rename_to()) in cc_op_gen_util.cc, which was unvalidated. Same pattern for the attr_name construction a few lines below it. Fix: expose IsValidAttrOrArgName() via op_def_util.h (previously file-local to op_def_util.cc) and validate rename_to() at its use sites in cc_op_gen_util.cc and python_op_gen.cc's ParamNames constructor, falling back to the original, already-validated OpDef name when rename_to() isn't a safe identifier. rename_to is a display-name preference, not required data, so a safe fallback is more appropriate than a hard failure. 3. Adds the suggested regression tests to op_def_util_test.cc, covering invalid characters and invalid start characters for both attribute and argument names. --- tensorflow/cc/framework/cc_op_gen_util.cc | 23 ++++++++++++--- tensorflow/core/framework/op_def_util.cc | 3 -- tensorflow/core/framework/op_def_util.h | 8 +++++ tensorflow/core/framework/op_def_util_test.cc | 29 +++++++++++++++++++ tensorflow/python/framework/python_op_gen.cc | 10 ++++++- 5 files changed, 65 insertions(+), 8 deletions(-) diff --git a/tensorflow/cc/framework/cc_op_gen_util.cc b/tensorflow/cc/framework/cc_op_gen_util.cc index 3186f29ef2794f..ba31f8899821a2 100644 --- a/tensorflow/cc/framework/cc_op_gen_util.cc +++ b/tensorflow/cc/framework/cc_op_gen_util.cc @@ -567,14 +567,22 @@ OpInfo::OpInfo(const OpDef& graph_op_def, const ApiDef& api_def, const auto& api_def_arg = *FindInputArg(api_def.arg_order(i), api_def); arg_types.push_back( absl::StrCat("::tensorflow::", ArgIsList(arg) ? "InputList" : "Input")); - arg_names.push_back(AvoidCPPKeywords(api_def_arg.rename_to())); + // rename_to() comes from ApiDef, a separate message from the OpDef + // arg/attr names validated by ValidateOpDef/ValidateArg. It is spliced + // as a raw C++ identifier (parameter name) below, so fall back to the + // already-validated original arg name if it isn't a safe identifier. + const std::string safe_input_name = + tensorflow::IsValidAttrOrArgName(api_def_arg.rename_to()) + ? api_def_arg.rename_to() + : arg.name(); + arg_names.push_back(AvoidCPPKeywords(safe_input_name)); // TODO(keveman): Include input type information. absl::string_view description = api_def_arg.description(); if (!description.empty()) { ConsumeEquals(&description); - absl::StrAppend(&comment, "* ", AvoidCPPKeywords(api_def_arg.rename_to()), - ": ", api_def_arg.description(), "\n"); + absl::StrAppend(&comment, "* ", AvoidCPPKeywords(safe_input_name), ": ", + api_def_arg.description(), "\n"); } } @@ -593,7 +601,14 @@ OpInfo::OpInfo(const OpDef& graph_op_def, const ApiDef& api_def, const auto entry = AttrTypeName(attr.type()); const auto attr_type_name = entry.first; const bool use_const = entry.second; - std::string attr_name = AvoidCPPKeywords(api_def_attr.rename_to()); + // See the safe_input_name comment above: rename_to() is unvalidated + // ApiDef data spliced as a raw C++ identifier here, so fall back to the + // already-validated original attr name if it isn't a safe identifier. + const std::string safe_attr_name = + tensorflow::IsValidAttrOrArgName(api_def_attr.rename_to()) + ? api_def_attr.rename_to() + : attr.name(); + std::string attr_name = AvoidCPPKeywords(safe_attr_name); std::string attr_comment; if (!api_def_attr.description().empty()) { diff --git a/tensorflow/core/framework/op_def_util.cc b/tensorflow/core/framework/op_def_util.cc index 64f82f89516cbe..61920cf7613a32 100644 --- a/tensorflow/core/framework/op_def_util.cc +++ b/tensorflow/core/framework/op_def_util.cc @@ -192,9 +192,6 @@ const ApiDef::Arg* FindInputArg(absl::string_view name, const ApiDef& api_def) { } \ } while (false) -// Forward declaration; defined below alongside IsValidOpName. -bool IsValidAttrOrArgName(absl::string_view sp); - static absl::Status ValidateArg(const OpDef::ArgDef& arg, const OpDef& op_def, bool output, absl::flat_hash_set* names) { diff --git a/tensorflow/core/framework/op_def_util.h b/tensorflow/core/framework/op_def_util.h index abaaeefb03c9a8..9504f05e6406e6 100644 --- a/tensorflow/core/framework/op_def_util.h +++ b/tensorflow/core/framework/op_def_util.h @@ -29,6 +29,14 @@ limitations under the License. namespace tensorflow { +// Returns true iff sp is a safe identifier: starts with a letter or +// underscore, and contains only letters, digits, and underscores. Used to +// validate any name (OpDef attr/arg names, or ApiDef rename_to values) +// that downstream code generators (cc_op_gen.cc, python_op_gen.cc) splice +// into generated C++/Python source, either as identifiers or into string +// literals. +bool IsValidAttrOrArgName(absl::string_view sp); + // Performs a consistency check across the fields of the op_def. absl::Status ValidateOpDef(const OpDef& op_def); diff --git a/tensorflow/core/framework/op_def_util_test.cc b/tensorflow/core/framework/op_def_util_test.cc index 41fd90d4e79fcf..7ad653e946db4d 100644 --- a/tensorflow/core/framework/op_def_util_test.cc +++ b/tensorflow/core/framework/op_def_util_test.cc @@ -68,6 +68,35 @@ void ExpectFailure(const absl::Status& status, const std::string& message) { } } // namespace +TEST_F(ValidateOpDefTest, InvalidAttrOrArgName) { + // Invalid characters in attribute name. + ExpectFailure( + TestProto( + "name: 'BadAttrName' attr { name: " + "'evil\\\"); system(\\\"touch /tmp/PWNED...\\\"); //' " + "type: 'int' }"), + "Invalid attr name"); + + // Invalid start character in attribute name. + ExpectFailure( + TestProto("name: 'BadAttrName' attr { name: '123_invalid' type: 'int' }"), + "Invalid attr name"); + + // Invalid characters in argument name. + ExpectFailure( + TestProto( + "name: 'BadArgName' input_arg { name: " + "'evil\\\"); system(\\\"touch /tmp/PWNED...\\\"); //' " + "type: DT_INT32 }"), + "Invalid argument name"); + + // Invalid start character in argument name. + ExpectFailure( + TestProto( + "name: 'BadArgName' input_arg { name: '123_invalid' type: DT_INT32 }"), + "Invalid argument name"); +} + TEST_F(ValidateOpDefTest, OpDefValid) { TF_EXPECT_OK(TestBuilder(OpDefBuilder("X").Attr("a: int"))); TF_EXPECT_OK(TestBuilder(OpDefBuilder("X").Input("a: int32"))); diff --git a/tensorflow/python/framework/python_op_gen.cc b/tensorflow/python/framework/python_op_gen.cc index 89eaa533e54d8c..e0263ad36c98ae 100644 --- a/tensorflow/python/framework/python_op_gen.cc +++ b/tensorflow/python/framework/python_op_gen.cc @@ -199,9 +199,17 @@ std::string DataTypeToPython(DataType dtype, const std::string& dtype_module); class ParamNames { public: // Create param based on Arg. + // + // rename_to comes from ApiDef, a separate message from the OpDef + // arg/attr names validated by ValidateOpDef/ValidateArg. It is spliced + // as a raw Python identifier (keyword argument name) below, so fall + // back to the already-validated original name if it isn't a safe + // identifier. ParamNames(const std::string& name, const std::string& rename_to) : name_(name) { - rename_to_ = AvoidPythonReserved(rename_to); + const std::string& safe_rename_to = + tensorflow::IsValidAttrOrArgName(rename_to) ? rename_to : name; + rename_to_ = AvoidPythonReserved(safe_rename_to); } // Get original parameter name. From 13a3adaa39ebb7e176b4df28c6d1d62badd8bd68 Mon Sep 17 00:00:00 2001 From: Cocoa Date: Fri, 21 Aug 2026 17:04:00 +0900 Subject: [PATCH 04/32] fix: compare strcmp against 0 when matching GPU delegate option keys Signed-off-by: Cocoa --- tensorflow/lite/delegates/gpu/delegate.cc | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/tensorflow/lite/delegates/gpu/delegate.cc b/tensorflow/lite/delegates/gpu/delegate.cc index 4e0c2767ca01d5..17a201013cc0a3 100644 --- a/tensorflow/lite/delegates/gpu/delegate.cc +++ b/tensorflow/lite/delegates/gpu/delegate.cc @@ -156,54 +156,54 @@ bool ParseOptions(const char* const* options_keys, const char* const* options_values, size_t num_options, TfLiteGpuDelegateOptionsV2* options) { for (size_t i = 0; i < num_options; ++i) { - if (strcmp(options_keys[i], "is_precision_loss_allowed")) { + if (strcmp(options_keys[i], "is_precision_loss_allowed") == 0) { if (!absl::SimpleAtoi(options_values[i], &options->is_precision_loss_allowed)) { TFLITE_LOG(TFLITE_LOG_WARNING, "ParseOptions: malformed option %s.", options_keys[i]); return false; } - } else if (strcmp(options_keys[i], "inference_preference")) { + } else if (strcmp(options_keys[i], "inference_preference") == 0) { if (!absl::SimpleAtoi(options_values[i], &options->inference_preference)) { TFLITE_LOG(TFLITE_LOG_WARNING, "ParseOptions: malformed option %s.", options_keys[i]); return false; } - } else if (strcmp(options_keys[i], "inference_priority1")) { + } else if (strcmp(options_keys[i], "inference_priority1") == 0) { if (!absl::SimpleAtoi(options_values[i], &options->inference_priority1)) { TFLITE_LOG(TFLITE_LOG_WARNING, "ParseOptions: malformed option %s.", options_keys[i]); return false; } - } else if (strcmp(options_keys[i], "inference_priority2")) { + } else if (strcmp(options_keys[i], "inference_priority2") == 0) { if (!absl::SimpleAtoi(options_values[i], &options->inference_priority2)) { TFLITE_LOG(TFLITE_LOG_WARNING, "ParseOptions: malformed option %s.", options_keys[i]); return false; } - } else if (strcmp(options_keys[i], "inference_priority3")) { + } else if (strcmp(options_keys[i], "inference_priority3") == 0) { if (!absl::SimpleAtoi(options_values[i], &options->inference_priority3)) { TFLITE_LOG(TFLITE_LOG_WARNING, "ParseOptions: malformed option %s.", options_keys[i]); return false; } - } else if (strcmp(options_keys[i], "experimental_flags")) { + } else if (strcmp(options_keys[i], "experimental_flags") == 0) { if (!absl::SimpleAtoi(options_values[i], &options->experimental_flags)) { TFLITE_LOG(TFLITE_LOG_WARNING, "ParseOptions: malformed option %s.", options_keys[i]); return false; } - } else if (strcmp(options_keys[i], "max_delegated_partitions")) { + } else if (strcmp(options_keys[i], "max_delegated_partitions") == 0) { if (!absl::SimpleAtoi(options_values[i], &options->max_delegated_partitions)) { TFLITE_LOG(TFLITE_LOG_WARNING, "ParseOptions: malformed option %s.", options_keys[i]); return false; } - } else if (strcmp(options_keys[i], "serialization_dir")) { + } else if (strcmp(options_keys[i], "serialization_dir") == 0) { options->serialization_dir = options_values[i]; - } else if (strcmp(options_keys[i], "model_token")) { + } else if (strcmp(options_keys[i], "model_token") == 0) { options->model_token = options_values[i]; } else { TFLITE_LOG(TFLITE_LOG_WARNING, "ParseOptions: unknown option %s.", From affb42de0e3db9d5a43a9e53fb0a8ca434d320e0 Mon Sep 17 00:00:00 2001 From: abhijeet117 Date: Sun, 23 Aug 2026 22:39:08 +0530 Subject: [PATCH 05/32] Use the IEEE sign bit in experimental.numpy.signbit signbit compared its argument with zero, so negative zero reported False even though NumPy reports True because the sign bit is set, and a negative NaN was also misreported. Bitcast float inputs to an integer type of the same size and compare that with zero instead. --- tensorflow/python/ops/numpy_ops/np_math_ops.py | 14 ++++++++++++++ .../python/ops/numpy_ops/np_math_ops_test.py | 14 ++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py index f35cb18eba0650..2c8b229cc0873d 100644 --- a/tensorflow/python/ops/numpy_ops/np_math_ops.py +++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py @@ -56,6 +56,15 @@ __name__, 'inf' ) +# Floating-point types mapped to an integer type of the same size, so their +# IEEE-754 sign bit can be checked with an integer comparison. +_SIGN_BITCAST_DTYPES = { + dtypes.bfloat16: dtypes.int16, + dtypes.float16: dtypes.int16, + dtypes.float32: dtypes.int32, + dtypes.float64: dtypes.int64, +} + @tf_export.tf_export('experimental.numpy.dot', v1=[]) @np_utils.np_doc_only('dot') @@ -809,6 +818,11 @@ def signbit(x): def f(x): if x.dtype == dtypes.bool: return array_ops.fill(array_ops.shape(x), False) + if x.dtype in _SIGN_BITCAST_DTYPES: + # Check the IEEE-754 sign bit instead of comparing with zero, which + # cannot tell -0.0 from +0.0 or a negative NaN from a positive one. + bits = array_ops.bitcast(x, _SIGN_BITCAST_DTYPES[x.dtype]) + return math_ops.less(bits, 0) return x < 0 return _scalar(f, x) diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py index f6c50d1e02babe..3dab0d29363a14 100644 --- a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py +++ b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py @@ -20,6 +20,7 @@ import numpy as np from tensorflow.python.eager import def_function +from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors from tensorflow.python.framework import ops from tensorflow.python.framework import tensor @@ -596,6 +597,19 @@ def testIsInf(self): self.assertFalse(np_math_ops.isneginf(x1)) self.assertFalse(np_math_ops.isneginf(x2)) + def testSignBit(self): + for transform in self.array_transforms: + values = transform([-1.5, -0., 0., 1.5]) + self.assertAllEqual( + np_math_ops.signbit(values), [True, True, False, False]) + # The sign bit is set even when the value compares equal to zero or NaN. + self.assertAllEqual( + np_math_ops.signbit([np.nan, -np.nan]), [False, True]) + self.assertAllEqual(np_math_ops.signbit([-3, 3]), [True, False]) + negative_zero = ops.convert_to_tensor( + [-0.], dtype=dtypes.bfloat16) + self.assertAllEqual(np_math_ops.signbit(negative_zero), [True]) + if __name__ == '__main__': tensor.enable_tensor_equality() ops.enable_eager_execution() From 311e2982429ac30a0be1a4b6c39ae7452a521611 Mon Sep 17 00:00:00 2001 From: prasanna8585 <65734642+prasanna8585@users.noreply.github.com> Date: Tue, 25 Aug 2026 20:20:30 +0530 Subject: [PATCH 06/32] fix: declare the absl/strings dependency op_def_util.cc actually uses Addresses dmiltr3's review: Google's internal presubmit on import failed because tensorflow/core/framework/BUILD's op_def_util cc_library target did not declare @com_google_absl//absl/strings in its deps, even though op_def_util.cc includes absl/strings/ascii.h (for absl::ascii_isalpha / absl::ascii_isalnum, used by this PR's own fix). Under strict layering checks this failed compilation for anything depending on op_def_util: error: module //third_party/tensorflow/core/framework:op_def_util does not directly depend on a module exporting 'third_party/absl/strings/ascii.h' ... Adds "@com_google_absl//absl/strings" to that target's deps, in the same place and form as the reviewer's own suggested diff. Verified: op_def_util.cc's use of absl/strings/ascii.h confirmed current before editing. Basic Starlark/Python-like syntax check on the whole BUILD file passes after the edit -- this sandbox has no Bazel available, so a real 'bazel build' could not be run here; please confirm CI passes. --- tensorflow/core/framework/BUILD | 1 + 1 file changed, 1 insertion(+) diff --git a/tensorflow/core/framework/BUILD b/tensorflow/core/framework/BUILD index 79f7b140907d4a..f37ae7ceb92dfa 100644 --- a/tensorflow/core/framework/BUILD +++ b/tensorflow/core/framework/BUILD @@ -1039,6 +1039,7 @@ cc_library( "//tensorflow/core/platform:types", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", + "@com_google_absl//absl/strings", ], ) From d5c83ecb5c206d3bb24bd39f47faffef4a91c4f0 Mon Sep 17 00:00:00 2001 From: PrabinDevkota Date: Wed, 26 Aug 2026 00:11:39 +0530 Subject: [PATCH 07/32] Add XLA kernels for MatrixDeterminant and LogMatrixDeterminant. --- .../compiler/jit/compilability_check_util.cc | 3 + .../compiler/jit/mark_for_compilation_pass.cc | 1 + tensorflow/compiler/tests/BUILD | 21 ++++ .../tests/matrix_determinant_op_test.py | 102 +++++++++++++++ tensorflow/compiler/tf2xla/kernels/BUILD | 20 +++ .../tf2xla/kernels/matrix_determinant_op.cc | 117 ++++++++++++++++++ 6 files changed, 264 insertions(+) create mode 100644 tensorflow/compiler/tests/matrix_determinant_op_test.py create mode 100644 tensorflow/compiler/tf2xla/kernels/matrix_determinant_op.cc diff --git a/tensorflow/compiler/jit/compilability_check_util.cc b/tensorflow/compiler/jit/compilability_check_util.cc index 8da8b2055c6c2b..c01514f6dcae1a 100644 --- a/tensorflow/compiler/jit/compilability_check_util.cc +++ b/tensorflow/compiler/jit/compilability_check_util.cc @@ -354,6 +354,7 @@ bool RecursiveCompilabilityChecker::OpIsInaccurate(const Node& node) const { bool RecursiveCompilabilityChecker::OpIsSlow(const Node& node) const { // b/128001705: SelfAdjointEigV2 and Svd performance issues. // b/135640736: MatrixInverse performance issues. + // MatrixDeterminant and LogMatrixDeterminant use the same QR path as Inverse. // b/111271662: MatrixSolve performance issues. // https://github.com/tensorflow/tensorflow/pull/31012: // ResizeNearestNeighbor, ResizeBilinear, and ResizeBilinearGrad sometimes @@ -364,6 +365,8 @@ bool RecursiveCompilabilityChecker::OpIsSlow(const Node& node) const { return node.type_string() == "SelfAdjointEigV2" || node.type_string() == "Svd" || node.type_string() == "Qr" || node.type_string() == "MatrixInverse" || + node.type_string() == "MatrixDeterminant" || + node.type_string() == "LogMatrixDeterminant" || node.type_string() == "MatrixSolve" || node.type_string() == "ResizeBilinearGrad" || node.type_string() == "NonMaxSuppressionV3" || diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass.cc b/tensorflow/compiler/jit/mark_for_compilation_pass.cc index de4f25150abb0d..72bc2d5f82fcfc 100644 --- a/tensorflow/compiler/jit/mark_for_compilation_pass.cc +++ b/tensorflow/compiler/jit/mark_for_compilation_pass.cc @@ -2152,6 +2152,7 @@ absl::flat_hash_set GetKnownXLAAllowlistOp() { "LowerBound", "MatMul", "MatrixBandPart", + "MatrixDeterminant", "MatrixDiag", "MatrixDiagPart", "MatrixDiagPartV2", diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD index 036599462e70d7..88defe8284bd47 100644 --- a/tensorflow/compiler/tests/BUILD +++ b/tensorflow/compiler/tests/BUILD @@ -172,6 +172,7 @@ tf_xla_combined_py_test( ], tests = [ # go/keep-sorted start + ":matrix_determinant_op_test_lib", ":matrix_inverse_op_test_lib", ":matrix_solve_op_test_lib", ":momentum_test_lib", @@ -603,6 +604,26 @@ tf_xla_py_strict_test( ], ) +tf_xla_py_strict_test( + name = "matrix_determinant_op_test", + size = "small", + timeout = "moderate", + srcs = ["matrix_determinant_op_test.py"], + tags = [ + "no_pip", # TODO(b/149738646): fix pip install so these tests run on kokoro pip + "notap", + ], + deps = [ + ":xla_test", + "//tensorflow/python/framework:dtypes", + "//tensorflow/python/ops:array_ops", + "//tensorflow/python/ops:linalg_ops", + "//tensorflow/python/ops:linalg_ops_gen", + "//tensorflow/python/platform:test", + "//third_party/py/numpy", + ], +) + tf_xla_py_strict_test( name = "matrix_inverse_op_test", size = "small", diff --git a/tensorflow/compiler/tests/matrix_determinant_op_test.py b/tensorflow/compiler/tests/matrix_determinant_op_test.py new file mode 100644 index 00000000000000..de8346422558e9 --- /dev/null +++ b/tensorflow/compiler/tests/matrix_determinant_op_test.py @@ -0,0 +1,102 @@ +# 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. +# ============================================================================== +"""Tests for XLA implementations of matrix determinant ops.""" + +import numpy as np + +from tensorflow.compiler.tests import xla_test +from tensorflow.python.framework import dtypes +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import gen_linalg_ops +from tensorflow.python.ops import linalg_ops +from tensorflow.python.platform import googletest + + +class DeterminantOpTest(xla_test.XLATestCase): + + def _verifyDeterminant(self, x, np_type): + y = x.astype(np_type) + if y.shape[-1] == 0 and y.shape[-2] == 0: + np_det = np.ones(y.shape[:-2], dtype=np_type) + np_sign = np.ones(y.shape[:-2], dtype=np_type) + np_log_abs = np.zeros(y.shape[:-2], dtype=np_type) + else: + np_det = np.array(np.linalg.det(y)).astype(np_type) + np_sign, np_log_abs = np.linalg.slogdet(y) + np_sign = np.array(np_sign).astype(np_type) + np_log_abs = np.array(np_log_abs).astype(np_type) + + with self.session() as sess: + p = array_ops.placeholder(dtypes.as_dtype(y.dtype), y.shape, name="x") + with self.test_scope(): + det = linalg_ops.matrix_determinant(p) + sign, log_abs = gen_linalg_ops.log_matrix_determinant(p) + det_out, sign_out, log_abs_out = sess.run( + [det, sign, log_abs], feed_dict={p: y}) + + self.assertAllClose(np_det, det_out, rtol=1e-3, atol=1e-3) + self.assertShapeEqual(np_det, det) + # Compare reconstructed determinants so a QR-vs-LU split of sign/log does + # not fail the test when the product still matches. Guard exp() so a + # singular matrix's -inf log-abs-det does not warn or overflow. + with np.errstate(over="ignore", invalid="ignore"): + np_recon = np_sign * np.exp(np_log_abs) + tf_recon = sign_out * np.exp(log_abs_out) + self.assertAllClose(np_recon, tf_recon, rtol=1e-3, atol=1e-3) + self.assertShapeEqual(np_sign, sign) + self.assertShapeEqual(np_log_abs, log_abs) + + def _verifyDeterminantReal(self, x): + for np_type in self.float_types & {np.float32, np.float64}: + self._verifyDeterminant(x, np_type) + + def testBasic(self): + # 1x1 + self._verifyDeterminantReal(np.array([[7.]])) + # 2x2 with negative determinant: det([[1, 2], [3, 4]]) == -2. + # This is the case xla::LogDet() gets wrong (returns NaN). + self._verifyDeterminantReal(np.array([[1., 2.], [3., 4.]])) + # 2x2 with positive determinant (the motivating jit_compile example). + self._verifyDeterminantReal(np.array([[4., 7.], [2., 6.]])) + # Singular. + self._verifyDeterminantReal(np.array([[0., 0.], [0., 0.]])) + # 3x3 with negative determinant. + self._verifyDeterminantReal( + np.array([[1., 0., 0.], [0., 1., 0.], [0., 0., -1.]])) + # Well-conditioned 5x5 triangular matrix; det = 2*3*4*5*6 = 720. + self._verifyDeterminantReal( + np.array([[2., 0., 0., 0., 0.], [1., 3., 0., 0., 0.], + [0., 1., 4., 0., 0.], [0., 0., 1., 5., 0.], + [0., 0., 0., 1., 6.]])) + + def testBatch(self): + # Mixed signs in the batch: dets are -2 and 3. + self._verifyDeterminantReal( + np.array([[[1., 2.], [3., 4.]], [[2., 1.], [1., 2.]]])) + matrix1 = np.array([[1., 2.], [3., 4.]]) + matrix2 = np.array([[1., 3.], [3., 5.]]) + batch = np.concatenate( + [np.expand_dims(matrix1, 0), + np.expand_dims(matrix2, 0)]) + batch = np.tile(batch, [2, 3, 1, 1]) + self._verifyDeterminantReal(batch) + + def testEmpty(self): + self._verifyDeterminantReal(np.empty([0, 2, 2])) + self._verifyDeterminantReal(np.empty([2, 0, 0])) + + +if __name__ == "__main__": + googletest.main() diff --git a/tensorflow/compiler/tf2xla/kernels/BUILD b/tensorflow/compiler/tf2xla/kernels/BUILD index 55a9669e361326..dd2ab20d4f31c9 100644 --- a/tensorflow/compiler/tf2xla/kernels/BUILD +++ b/tensorflow/compiler/tf2xla/kernels/BUILD @@ -112,6 +112,7 @@ cc_library( ":lrn_ops", ":matmul_op", ":matrix_band_part_op", + ":matrix_determinant_op", ":matrix_diag_ops", ":matrix_inverse_op", ":matrix_solve_op", @@ -1026,6 +1027,25 @@ tf_kernel_library( ], ) +tf_kernel_library( + name = "matrix_determinant_op", + srcs = ["matrix_determinant_op.cc"], + deps = [ + "//tensorflow/compiler/tf2xla:xla_compilation_device", + "//tensorflow/compiler/tf2xla:xla_compiler", + "//tensorflow/compiler/tf2xla:xla_context", + "//tensorflow/compiler/tf2xla:xla_op_registry", + "//tensorflow/compiler/tf2xla:xla_resource", + "//tensorflow/compiler/tf2xla/ops:xla_ops", + "//tensorflow/core:framework", + "@com_google_absl//absl/status", + "@com_google_absl//absl/strings", + "@xla//xla/hlo/builder:xla_builder", + "@xla//xla/hlo/builder/lib:constants", + "@xla//xla/hlo/builder/lib:logdet", + ], +) + tf_kernel_library( name = "matrix_inverse_op", srcs = ["matrix_inverse_op.cc"], diff --git a/tensorflow/compiler/tf2xla/kernels/matrix_determinant_op.cc b/tensorflow/compiler/tf2xla/kernels/matrix_determinant_op.cc new file mode 100644 index 00000000000000..1bd66e0d7ad702 --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/matrix_determinant_op.cc @@ -0,0 +1,117 @@ +/* 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 "absl/status/status.h" +#include "absl/strings/str_cat.h" +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/op_requires.h" +#include "tensorflow/core/framework/tensor_shape.h" +#include "xla/hlo/builder/lib/constants.h" +#include "xla/hlo/builder/lib/logdet.h" +#include "xla/hlo/builder/xla_builder.h" + +namespace tensorflow { +namespace { + +absl::Status CheckSquareMatrix(const TensorShape& input_shape) { + const int64_t ndims = input_shape.dims(); + if (ndims < 2) { + return absl::InvalidArgumentError( + absl::StrCat("Input must have rank >= 2, got ", ndims)); + } + if (input_shape.dim_size(ndims - 2) != input_shape.dim_size(ndims - 1)) { + return absl::InvalidArgumentError(absl::StrCat( + "Input matrices must be square, got ", input_shape.dim_size(ndims - 2), + " != ", input_shape.dim_size(ndims - 1))); + } + return absl::OkStatus(); +} + +// Broadcasts a scalar to the batch shape of a [..., n, n] matrix input. +xla::XlaOp BroadcastScalarToBatch(xla::XlaOp scalar, + const TensorShape& input_shape) { + std::vector batch_dims(input_shape.dims() - 2); + for (int i = 0; i < input_shape.dims() - 2; ++i) { + batch_dims[i] = input_shape.dim_size(i); + } + return xla::Broadcast(scalar, batch_dims); +} + +// slogdet(A) = (sign, log|det|). For n == 0, det is defined to be 1. +// xla::SLogDet cannot handle that case: it slices Householder taus to n-1. +xla::SignAndLogDet SLogDetOrEmpty(XlaOpKernelContext* ctx, + const TensorShape& input_shape) { + const int64_t n = input_shape.dim_size(input_shape.dims() - 1); + if (n == 0) { + const xla::PrimitiveType type = ctx->input_xla_type(0); + return xla::SignAndLogDet{ + BroadcastScalarToBatch(xla::One(ctx->builder(), type), input_shape), + BroadcastScalarToBatch(xla::Zero(ctx->builder(), type), input_shape)}; + } + return xla::SLogDet(ctx->Input(0)); +} + +class MatrixDeterminantOp : public XlaOpKernel { + public: + explicit MatrixDeterminantOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} + + void Compile(XlaOpKernelContext* ctx) override { + const TensorShape input_shape = ctx->InputShape(0); + OP_REQUIRES_OK(ctx, CheckSquareMatrix(input_shape)); + + // det = sign * exp(log|det|). Do not use xla::LogDet(): that returns NaN + // for matrices with a negative determinant. + const xla::SignAndLogDet slogdet = SLogDetOrEmpty(ctx, input_shape); + ctx->SetOutput(0, slogdet.sign * xla::Exp(slogdet.logdet)); + } + + private: + MatrixDeterminantOp(const MatrixDeterminantOp&) = delete; + void operator=(const MatrixDeterminantOp&) = delete; +}; + +class LogMatrixDeterminantOp : public XlaOpKernel { + public: + explicit LogMatrixDeterminantOp(OpKernelConstruction* ctx) + : XlaOpKernel(ctx) {} + + void Compile(XlaOpKernelContext* ctx) override { + const TensorShape input_shape = ctx->InputShape(0); + OP_REQUIRES_OK(ctx, CheckSquareMatrix(input_shape)); + + const xla::SignAndLogDet slogdet = SLogDetOrEmpty(ctx, input_shape); + ctx->SetOutput(0, slogdet.sign); + ctx->SetOutput(1, slogdet.logdet); + } + + private: + LogMatrixDeterminantOp(const LogMatrixDeterminantOp&) = delete; + void operator=(const LogMatrixDeterminantOp&) = delete; +}; + +// TODO(b/135640736): Allow complex types once XLA QR/SLogDet is validated for +// them, matching MatrixInverse. +REGISTER_XLA_OP(Name("MatrixDeterminant").TypeConstraint("T", kFloatTypes), + MatrixDeterminantOp); +REGISTER_XLA_OP(Name("LogMatrixDeterminant").TypeConstraint("T", kFloatTypes), + LogMatrixDeterminantOp); + +} // namespace +} // namespace tensorflow From aea3853658a2f3a5e27ab223121acfb41dcac255 Mon Sep 17 00:00:00 2001 From: PrabinDevkota Date: Wed, 26 Aug 2026 00:21:04 +0530 Subject: [PATCH 08/32] Add a [0, 0] empty-matrix case to the XLA determinant test. --- tensorflow/compiler/tests/matrix_determinant_op_test.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tensorflow/compiler/tests/matrix_determinant_op_test.py b/tensorflow/compiler/tests/matrix_determinant_op_test.py index de8346422558e9..f237eecb83f830 100644 --- a/tensorflow/compiler/tests/matrix_determinant_op_test.py +++ b/tensorflow/compiler/tests/matrix_determinant_op_test.py @@ -94,6 +94,7 @@ def testBatch(self): self._verifyDeterminantReal(batch) def testEmpty(self): + self._verifyDeterminantReal(np.empty([0, 0])) self._verifyDeterminantReal(np.empty([0, 2, 2])) self._verifyDeterminantReal(np.empty([2, 0, 0])) From 5490c0f42dad72b06e8c92a5c4a130757db216e6 Mon Sep 17 00:00:00 2001 From: prasanna8585 <65734642+prasanna8585@users.noreply.github.com> Date: Wed, 26 Aug 2026 01:21:05 +0530 Subject: [PATCH 09/32] fix: move GHSA-2xp6-8g4h-qw72's defense to the code-gen splice sites Addresses review: the runtime validation added to ValidateArg/ ValidateOpDef broke internal presubmit across 1,400+ targets. Root cause: several legitimate, already-registered ops (e.g. TFLite_Detection_PostProcess, registered in tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc) use argument names containing '/' and ':' ("raw_outputs/box_encodings", "TFLite_Detection_PostProcess:1") for reasons unrelated to code generation. Rejecting those at OpDef registration/validation time -- which runs for every op in the runtime, not just ones that go through code generation -- broke that legitimate, pre-existing usage. The actual injection risk only exists where a name is spliced, unescaped, as a raw identifier into generated C++/Python source by cc_op_gen.cc / cc_op_gen_util.cc / python_op_gen.cc. Moves the defense to exactly those splice sites instead of the runtime-wide check: 1. Removes the IsValidAttrOrArgName VALIDATE() calls from ValidateArg and ValidateOpDef, keeping IsValidAttrOrArgName itself as a utility function for code generators. 2. Both cc_op_gen_util.cc and python_op_gen.cc already had a fallback pattern for ApiDef rename_to() built on the assumption that the ORIGINAL arg/attr name was "already validated" by the runtime check removed in (1). That assumption no longer holds, so all three call sites are upgraded to a three-tier fallback: rename_to if safe, else the original name if THAT is safe, else a new SanitizeToIdentifier() helper (added alongside IsValidAttrOrArgName in op_def_util.{h,cc}) that guarantees a syntactically safe identifier for any input by replacing unsafe characters with underscores. 3. Replaces op_def_util_test.cc's InvalidAttrOrArgName test (which asserted ValidateOpDef rejects non-identifier names) with direct unit tests of IsValidAttrOrArgName and SanitizeToIdentifier as pure functions, plus a positive regression test confirming ValidateOpDef now accepts TFLite_Detection_PostProcess's exact real-world argument names. Verified with a standalone C++ harness (no Bazel available in this environment) reproducing both functions exactly: 14 checks, including confirming every SanitizeToIdentifier output itself re-passes IsValidAttrOrArgName -- the invariant the whole three-tier fallback depends on. All 14 pass. The exact GHSA-2xp6-8g4h-qw72 payload and TFLite_Detection_PostProcess's real argument names are both correctly rejected by IsValidAttrOrArgName and correctly sanitized into safe identifiers by SanitizeToIdentifier. Please confirm the real build and the 1,400+ previously-failing presubmit targets pass in CI before merging -- this environment has no Bazel to verify that directly. --- tensorflow/cc/framework/cc_op_gen_util.cc | 26 +++++-- tensorflow/core/framework/op_def_util.cc | 58 ++++++++++---- tensorflow/core/framework/op_def_util.h | 9 ++- tensorflow/core/framework/op_def_util_test.cc | 75 +++++++++++++------ tensorflow/python/framework/python_op_gen.cc | 16 ++-- 5 files changed, 132 insertions(+), 52 deletions(-) diff --git a/tensorflow/cc/framework/cc_op_gen_util.cc b/tensorflow/cc/framework/cc_op_gen_util.cc index ba31f8899821a2..3560b5f70705ab 100644 --- a/tensorflow/cc/framework/cc_op_gen_util.cc +++ b/tensorflow/cc/framework/cc_op_gen_util.cc @@ -568,13 +568,19 @@ OpInfo::OpInfo(const OpDef& graph_op_def, const ApiDef& api_def, arg_types.push_back( absl::StrCat("::tensorflow::", ArgIsList(arg) ? "InputList" : "Input")); // rename_to() comes from ApiDef, a separate message from the OpDef - // arg/attr names validated by ValidateOpDef/ValidateArg. It is spliced - // as a raw C++ identifier (parameter name) below, so fall back to the - // already-validated original arg name if it isn't a safe identifier. + // arg/attr names. Both are spliced as a raw C++ identifier (parameter + // name) below: prefer rename_to() when it's a safe identifier, fall + // back to the original arg name when THAT is safe, and only sanitize + // as a last resort, since IsValidAttrOrArgName is not enforced at + // OpDef registration and a legitimately-registered op's argument name + // is not guaranteed to already be a safe identifier (e.g. + // TFLite_Detection_PostProcess's "raw_outputs/box_encodings"). const std::string safe_input_name = tensorflow::IsValidAttrOrArgName(api_def_arg.rename_to()) ? api_def_arg.rename_to() - : arg.name(); + : tensorflow::IsValidAttrOrArgName(arg.name()) + ? arg.name() + : tensorflow::SanitizeToIdentifier(arg.name()); arg_names.push_back(AvoidCPPKeywords(safe_input_name)); // TODO(keveman): Include input type information. @@ -601,13 +607,17 @@ OpInfo::OpInfo(const OpDef& graph_op_def, const ApiDef& api_def, const auto entry = AttrTypeName(attr.type()); const auto attr_type_name = entry.first; const bool use_const = entry.second; - // See the safe_input_name comment above: rename_to() is unvalidated - // ApiDef data spliced as a raw C++ identifier here, so fall back to the - // already-validated original attr name if it isn't a safe identifier. + // See the safe_input_name comment above: rename_to() and the original + // attr name are both unvalidated at OpDef-registration time and both + // get spliced as a raw C++ identifier here, so prefer rename_to() when + // safe, fall back to the original name when THAT is safe, and only + // sanitize as a last resort. const std::string safe_attr_name = tensorflow::IsValidAttrOrArgName(api_def_attr.rename_to()) ? api_def_attr.rename_to() - : attr.name(); + : tensorflow::IsValidAttrOrArgName(attr.name()) + ? attr.name() + : tensorflow::SanitizeToIdentifier(attr.name()); std::string attr_name = AvoidCPPKeywords(safe_attr_name); std::string attr_comment; diff --git a/tensorflow/core/framework/op_def_util.cc b/tensorflow/core/framework/op_def_util.cc index 61920cf7613a32..367e7a943b11db 100644 --- a/tensorflow/core/framework/op_def_util.cc +++ b/tensorflow/core/framework/op_def_util.cc @@ -198,10 +198,6 @@ static absl::Status ValidateArg(const OpDef::ArgDef& arg, const OpDef& op_def, const std::string suffix = absl::StrCat(output ? " for output '" : " for input '", arg.name(), "'"); VALIDATE(names->emplace(arg.name()).second, "Duplicate name: ", arg.name()); - VALIDATE(IsValidAttrOrArgName(arg.name()), "Invalid argument name: ", - arg.name(), - " (must start with a letter or underscore and contain only " - "letters, digits, and underscores)"); VALIDATE(HasAttrStyleType(arg), "Missing type", suffix); if (!arg.number_attr().empty()) { @@ -273,12 +269,15 @@ bool IsValidOpName(absl::string_view sp) { } // Attribute and argument names are not restricted to CamelCase like op -// names, but they are spliced verbatim (unescaped) into generated C++ and -// Python wrapper source by tools such as cc_op_gen.cc and python_op_gen.cc. -// To prevent a malicious attr/arg name from injecting arbitrary code into -// that generated source (e.g. a name containing `", ), ;`), restrict these -// names to a safe identifier character set: they must start with a letter -// or underscore and contain only letters, digits, and underscores. +// names. Some code generators, such as cc_op_gen.cc and python_op_gen.cc, +// splice a name verbatim (unescaped) into generated C++ or Python wrapper +// source as a raw identifier; this function tells such a call site whether +// a given name is safe to use that way as-is. It is not, and must not be +// used as, a runtime validity check on OpDef registration: legitimate, +// already-registered ops can and do use argument names outside this safe +// set for reasons unrelated to code generation (see SanitizeToIdentifier's +// comment below for a concrete example), and rejecting them at +// registration time breaks that legitimate usage. bool IsValidAttrOrArgName(absl::string_view sp) { if (sp.empty()) return false; if (sp[0] != '_' && !absl::ascii_isalpha(sp[0])) { @@ -293,6 +292,41 @@ bool IsValidAttrOrArgName(absl::string_view sp) { return true; } +// Converts `sp` into a safe identifier by replacing any character outside +// IsValidAttrOrArgName's safe set with an underscore, and prefixing with an +// underscore if the result would otherwise start with a digit or be empty. +// +// Some legitimate, already-registered OpDefs use argument names outside +// that safe set for reasons unrelated to code generation -- for example +// TFLite_Detection_PostProcess registers inputs like +// "raw_outputs/box_encodings" and outputs like +// "TFLite_Detection_PostProcess:1". Rejecting such an OpDef at +// registration/validation time (as ValidateOpDef/ValidateArg once did) +// breaks that legitimate usage; the actual splice-into-generated-source +// risk IsValidAttrOrArgName exists for only applies where a name is used +// as a raw identifier by a code generator such as cc_op_gen.cc or +// python_op_gen.cc. This function gives those call sites a fallback that +// is always a syntactically safe identifier, so they can sanitize instead +// of rejecting -- callers should first check IsValidAttrOrArgName and use +// the name as-is when it already passes, calling this only for a name +// that doesn't. +std::string SanitizeToIdentifier(absl::string_view sp) { + std::string result; + result.reserve(sp.size() + 1); + for (char c : sp) { + if (c == '_' || absl::ascii_isalnum(c)) { + result.push_back(c); + } else { + result.push_back('_'); + } + } + if (result.empty() || + (!absl::ascii_isalpha(result[0]) && result[0] != '_')) { + result.insert(result.begin(), '_'); + } + return result; +} + absl::Status ValidateOpDef(const OpDef& op_def) { if (!absl::StartsWith(op_def.name(), "_")) { VALIDATE(IsValidOpName(op_def.name()), "Invalid name: ", op_def.name(), @@ -305,10 +339,6 @@ absl::Status ValidateOpDef(const OpDef& op_def) { // Validate name VALIDATE(names.emplace(attr.name()).second, "Duplicate name: ", attr.name()); - VALIDATE(IsValidAttrOrArgName(attr.name()), "Invalid attr name: ", - attr.name(), - " (must start with a letter or underscore and contain only " - "letters, digits, and underscores)"); DataType dt; VALIDATE(!DataTypeFromString(attr.name(), &dt), "Attr can't have name ", attr.name(), " that matches a data type"); diff --git a/tensorflow/core/framework/op_def_util.h b/tensorflow/core/framework/op_def_util.h index 9504f05e6406e6..7bf554a21a934f 100644 --- a/tensorflow/core/framework/op_def_util.h +++ b/tensorflow/core/framework/op_def_util.h @@ -34,9 +34,16 @@ namespace tensorflow { // validate any name (OpDef attr/arg names, or ApiDef rename_to values) // that downstream code generators (cc_op_gen.cc, python_op_gen.cc) splice // into generated C++/Python source, either as identifiers or into string -// literals. +// literals. Not a runtime validity check on OpDef registration -- see the +// definition in op_def_util.cc for why. bool IsValidAttrOrArgName(absl::string_view sp); +// Converts `sp` into a safe identifier (see IsValidAttrOrArgName) by +// replacing unsafe characters with underscores. A code generator that +// needs a raw identifier from a name that fails IsValidAttrOrArgName +// should use this rather than the name itself. +std::string SanitizeToIdentifier(absl::string_view sp); + // Performs a consistency check across the fields of the op_def. absl::Status ValidateOpDef(const OpDef& op_def); diff --git a/tensorflow/core/framework/op_def_util_test.cc b/tensorflow/core/framework/op_def_util_test.cc index 7ad653e946db4d..df9b64bd86ee61 100644 --- a/tensorflow/core/framework/op_def_util_test.cc +++ b/tensorflow/core/framework/op_def_util_test.cc @@ -68,33 +68,60 @@ void ExpectFailure(const absl::Status& status, const std::string& message) { } } // namespace -TEST_F(ValidateOpDefTest, InvalidAttrOrArgName) { - // Invalid characters in attribute name. - ExpectFailure( - TestProto( - "name: 'BadAttrName' attr { name: " - "'evil\\\"); system(\\\"touch /tmp/PWNED...\\\"); //' " - "type: 'int' }"), - "Invalid attr name"); +// IsValidAttrOrArgName and SanitizeToIdentifier are code-generator-facing +// utilities, not a runtime validity constraint on OpDef registration -- +// ValidateOpDef/ValidateArg deliberately do not call them (see +// OpDefValidAcceptsNonIdentifierNames below and the comment on +// IsValidAttrOrArgName's definition for why). +TEST(IsValidAttrOrArgNameTest, AcceptsSafeIdentifiers) { + EXPECT_TRUE(IsValidAttrOrArgName("dtype")); + EXPECT_TRUE(IsValidAttrOrArgName("T")); + EXPECT_TRUE(IsValidAttrOrArgName("num_threads_2")); + EXPECT_TRUE(IsValidAttrOrArgName("_internal")); +} - // Invalid start character in attribute name. - ExpectFailure( - TestProto("name: 'BadAttrName' attr { name: '123_invalid' type: 'int' }"), - "Invalid attr name"); +TEST(IsValidAttrOrArgNameTest, RejectsUnsafeNames) { + // The exact malicious payload from GHSA-2xp6-8g4h-qw72. + EXPECT_FALSE( + IsValidAttrOrArgName("evil\"); system(\"touch /tmp/PWNED\"); //")); + // A legitimate, already-registered op's argument name that is not a + // safe identifier for reasons unrelated to code generation. + EXPECT_FALSE(IsValidAttrOrArgName("raw_outputs/box_encodings")); + EXPECT_FALSE(IsValidAttrOrArgName("TFLite_Detection_PostProcess:1")); + EXPECT_FALSE(IsValidAttrOrArgName("")); + EXPECT_FALSE(IsValidAttrOrArgName("1abc")); +} - // Invalid characters in argument name. - ExpectFailure( - TestProto( - "name: 'BadArgName' input_arg { name: " - "'evil\\\"); system(\\\"touch /tmp/PWNED...\\\"); //' " - "type: DT_INT32 }"), - "Invalid argument name"); +TEST(SanitizeToIdentifierTest, ProducesASafeIdentifierForAnyInput) { + // Every unsafe character becomes an underscore; the malicious payload + // that IsValidAttrOrArgName exists to reject becomes an inert, safe + // identifier string rather than being spliced verbatim. + EXPECT_EQ(SanitizeToIdentifier("evil\"); system(\"touch /tmp/PWNED\"); //"), + "evil____system__touch__tmp_PWNED______"); + EXPECT_EQ(SanitizeToIdentifier("raw_outputs/box_encodings"), + "raw_outputs_box_encodings"); + EXPECT_EQ(SanitizeToIdentifier("TFLite_Detection_PostProcess:1"), + "TFLite_Detection_PostProcess_1"); + // A result that would start with a digit gets a leading underscore + // rather than becoming a syntactically invalid identifier. + EXPECT_EQ(SanitizeToIdentifier("123_invalid"), "_123_invalid"); + EXPECT_EQ(SanitizeToIdentifier(""), "_"); + // A name that is already safe is unaffected other than round-tripping + // through the same character-by-character pass. + EXPECT_EQ(SanitizeToIdentifier("already_safe"), "already_safe"); +} - // Invalid start character in argument name. - ExpectFailure( - TestProto( - "name: 'BadArgName' input_arg { name: '123_invalid' type: DT_INT32 }"), - "Invalid argument name"); +// The regression this fix closes: ValidateOpDef must accept an OpDef using +// argument names outside IsValidAttrOrArgName's safe set, matching a real, +// already-registered op (TFLite_Detection_PostProcess, registered in +// tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc) that +// broke under the runtime rejection this fix removes. +TEST_F(ValidateOpDefTest, OpDefValidAcceptsNonIdentifierNames) { + TF_EXPECT_OK(TestProto( + "name: 'TFLite_Detection_PostProcess' " + "input_arg { name: 'raw_outputs/box_encodings' type: DT_FLOAT } " + "input_arg { name: 'raw_outputs/class_predictions' type: DT_FLOAT } " + "output_arg { name: 'TFLite_Detection_PostProcess:1' type: DT_FLOAT }")); } TEST_F(ValidateOpDefTest, OpDefValid) { diff --git a/tensorflow/python/framework/python_op_gen.cc b/tensorflow/python/framework/python_op_gen.cc index e0263ad36c98ae..b713871cbed36a 100644 --- a/tensorflow/python/framework/python_op_gen.cc +++ b/tensorflow/python/framework/python_op_gen.cc @@ -201,14 +201,20 @@ class ParamNames { // Create param based on Arg. // // rename_to comes from ApiDef, a separate message from the OpDef - // arg/attr names validated by ValidateOpDef/ValidateArg. It is spliced - // as a raw Python identifier (keyword argument name) below, so fall - // back to the already-validated original name if it isn't a safe - // identifier. + // arg/attr names. Both are spliced as a raw Python identifier (keyword + // argument name) below: prefer rename_to when it's a safe identifier, + // fall back to the original name when THAT is safe, and only sanitize + // as a last resort, since IsValidAttrOrArgName is not enforced at OpDef + // registration and a legitimately-registered op's argument name is not + // guaranteed to already be a safe identifier (e.g. + // TFLite_Detection_PostProcess's "raw_outputs/box_encodings"). ParamNames(const std::string& name, const std::string& rename_to) : name_(name) { const std::string& safe_rename_to = - tensorflow::IsValidAttrOrArgName(rename_to) ? rename_to : name; + tensorflow::IsValidAttrOrArgName(rename_to) ? rename_to + : tensorflow::IsValidAttrOrArgName(name) + ? name + : tensorflow::SanitizeToIdentifier(name); rename_to_ = AvoidPythonReserved(safe_rename_to); } From 2b415a7d25be77ce15046a07632d2b88fdea69fe Mon Sep 17 00:00:00 2001 From: prasanna8585 <65734642+prasanna8585@users.noreply.github.com> Date: Wed, 26 Aug 2026 09:37:44 +0530 Subject: [PATCH 10/32] fix: sanitize the remaining unvalidated rename_to() splice sites in the C++ op generators Addresses review round 3 (SensitiveApiFuzzingCheck): several C++ operator-wrapper code-generation paths in cc_op_gen_util.cc and cc_op_gen.cc still concatenated raw, unsanitized ApiDef rename_to() values directly into generated C++ source, even after the previous round fixed python_op_gen.cc and two call sites in cc_op_gen_util.cc. Adds a single shared SafeRenameTo(name, rename_to) helper to cc_op_gen_util.{h,cc} (the shared C++ code-gen utility file both cc_op_gen.cc and cc_op_gen_util.cc already include), matching the review's exact suggested signature and three-tier fallback: rename_to if it's a safe identifier, else the original name if THAT is safe, else SanitizeToIdentifier(name). Refactors the two previously-fixed call sites to use this shared helper too, replacing duplicated inline ternary chains with one implementation. Fixes every remaining unsanitized splice site found, beyond the reviewer's own list: - cc_op_gen_util.cc: output_names (mirrors the already-fixed input case in the same constructor) - cc_op_gen_util.cc's GetOpAttrStruct(): four separate splice points for the same attribute (the setter method name, the field-access expression, the static defaults function name, and the struct field declaration) -- computed once per loop iteration so all four agree on the same identifier - cc_op_gen.cc: the static Attrs::(x) helper generator (a near-duplicate of GetOpAttrStruct's setter-name generation) - cc_op_gen.cc's GetConstructorBody(): the input-arg declaration loop AND the separate .Input() reference loop, which reference the same generated local variable by name from two different loops over parallel ApiDef arg lists -- both fixed with the identical sanitization call so the generated .Input(_name) reference always matches its auto _name = ... declaration - cc_op_gen.cc's .Attr() value-reference (both the attrs._ field-access form and the bare-identifier form) Verified with standalone C++ harnesses (no Bazel in this environment): SafeRenameTo itself passes 6 checks (normal case, unsafe-rename_to/ safe-name fallback, TFLite_Detection_PostProcess's real both-unsafe case, the exact GHSA payload in both fields, empty rename_to, a colon-suffixed output name), confirming every result is itself a valid identifier. A second harness specifically confirms the GetConstructorBody consistency invariant: the declaration loop and the .Input() loop, called independently with the same (name, rename_to) pair, always produce byte-identical output. Confirmed via grep that zero unsanitized .rename_to() calls remain in either file -- every occurrence is now an argument to SafeRenameTo. Please confirm the real build passes (specifically //tensorflow/cc:cc_ops or equivalent) before merging. --- tensorflow/cc/framework/cc_op_gen.cc | 36 +++++++++++++++--- tensorflow/cc/framework/cc_op_gen_util.cc | 46 +++++++++++++++-------- tensorflow/cc/framework/cc_op_gen_util.h | 11 ++++++ 3 files changed, 71 insertions(+), 22 deletions(-) diff --git a/tensorflow/cc/framework/cc_op_gen.cc b/tensorflow/cc/framework/cc_op_gen.cc index fb4e752c9b9533..baadff11254ec5 100644 --- a/tensorflow/cc/framework/cc_op_gen.cc +++ b/tensorflow/cc/framework/cc_op_gen.cc @@ -119,7 +119,12 @@ void WriteClassDecl(const OpInfo& op_info, WritableFile* h) { const auto entry = AttrTypeName(attr.type()); const auto attr_type_name = entry.first; const bool use_const = entry.second; - const std::string camel_case_name = ToCamelCase(api_def_attr.rename_to()); + // See cc_op_gen_util.cc's GetOpAttrStruct: attr.name() and + // api_def_attr.rename_to() are both unvalidated at OpDef-registration + // time and both spliced as a raw C++ identifier below. + const std::string safe_attr_name = + SafeRenameTo(attr.name(), api_def_attr.rename_to()); + const std::string camel_case_name = ToCamelCase(safe_attr_name); const std::string suffix = (camel_case_name == op_info.op_name || camel_case_name == "Attrs") ? "_" @@ -212,10 +217,19 @@ std::string GetConstructorBody(const OpInfo& op_info) { for (int i = 0; i < op_info.graph_op_def.input_arg_size(); ++i) { const auto& arg(op_info.graph_op_def.input_arg(i)); const auto& api_def_arg(op_info.api_def.in_arg(i)); + // See cc_op_gen_util.cc's GetOpAttrStruct comment: arg.name() and + // api_def_arg.rename_to() are both unvalidated at OpDef-registration + // time. Spliced twice below -- once as the local variable name + // (`_`), once via AvoidCPPKeywords as the argument reference -- + // so both must agree with each other and with the corresponding + // .Input(_) reference generated further down from the same + // (arg.name(), rename_to()) pair. + const std::string safe_arg_name = + SafeRenameTo(arg.name(), api_def_arg.rename_to()); strings::StrAppend( - &body, " auto _", api_def_arg.rename_to(), " = ::tensorflow::ops::", + &body, " auto _", safe_arg_name, " = ::tensorflow::ops::", ArgIsList(arg) ? "AsNodeOutList" : "AsNodeOut", "(", scope_str, ", ", - AvoidCPPKeywords(api_def_arg.rename_to()), ");\n"); + AvoidCPPKeywords(safe_arg_name), ");\n"); absl::StrAppend(&body, " ", return_on_error, "\n"); } @@ -228,7 +242,11 @@ std::string GetConstructorBody(const OpInfo& op_info) { const std::string spaces = " "; for (int i = 0; i < op_info.api_def.in_arg_size(); ++i) { const auto& arg(op_info.api_def.in_arg(i)); - absl::StrAppend(&body, spaces, ".Input(_", arg.rename_to(), ")\n"); + // Same (name, rename_to) pair as the declaration loop above, for the + // same index -- SafeRenameTo is pure, so this independently produces + // the identical identifier the declaration above already emitted. + const std::string safe_arg_name = SafeRenameTo(arg.name(), arg.rename_to()); + absl::StrAppend(&body, spaces, ".Input(_", safe_arg_name, ")\n"); } for (int i = 0; i < op_info.api_def.attr_size(); ++i) { const auto& graph_attr(op_info.graph_op_def.attr(i)); @@ -237,10 +255,15 @@ std::string GetConstructorBody(const OpInfo& op_info) { op_info.inferred_input_attrs.end()) { continue; } + // See above: graph_attr.name() and api_def_attr.rename_to() are both + // unvalidated, and both are spliced as a raw C++ identifier/field + // access below. + const std::string safe_attr_name = + SafeRenameTo(graph_attr.name(), api_def_attr.rename_to()); const std::string attr_name = api_def_attr.has_default_value() - ? absl::StrCat("attrs.", api_def_attr.rename_to(), "_") - : AvoidCPPKeywords(api_def_attr.rename_to()); + ? absl::StrCat("attrs.", safe_attr_name, "_") + : AvoidCPPKeywords(safe_attr_name); strings::StrAppend(&body, spaces, ".Attr(\"", absl::CEscape(graph_attr.name()), "\", ", attr_name, ")\n"); @@ -445,3 +468,4 @@ void WriteCCOps(const OpList& ops, const ApiDefMap& api_def_map, } // namespace cc_op } // namespace tensorflow + diff --git a/tensorflow/cc/framework/cc_op_gen_util.cc b/tensorflow/cc/framework/cc_op_gen_util.cc index 3560b5f70705ab..2d909afb0100a8 100644 --- a/tensorflow/cc/framework/cc_op_gen_util.cc +++ b/tensorflow/cc/framework/cc_op_gen_util.cc @@ -495,6 +495,16 @@ std::string AvoidCPPKeywords(absl::string_view name) { return std::string(name); } +std::string SafeRenameTo(absl::string_view name, absl::string_view rename_to) { + if (tensorflow::IsValidAttrOrArgName(rename_to)) { + return std::string(rename_to); + } + if (tensorflow::IsValidAttrOrArgName(name)) { + return std::string(name); + } + return tensorflow::SanitizeToIdentifier(name); +} + void InferArgAttributes( const OpDef::ArgDef& arg, std::unordered_map* inferred_attrs) { @@ -576,11 +586,7 @@ OpInfo::OpInfo(const OpDef& graph_op_def, const ApiDef& api_def, // is not guaranteed to already be a safe identifier (e.g. // TFLite_Detection_PostProcess's "raw_outputs/box_encodings"). const std::string safe_input_name = - tensorflow::IsValidAttrOrArgName(api_def_arg.rename_to()) - ? api_def_arg.rename_to() - : tensorflow::IsValidAttrOrArgName(arg.name()) - ? arg.name() - : tensorflow::SanitizeToIdentifier(arg.name()); + SafeRenameTo(arg.name(), api_def_arg.rename_to()); arg_names.push_back(AvoidCPPKeywords(safe_input_name)); // TODO(keveman): Include input type information. @@ -613,11 +619,7 @@ OpInfo::OpInfo(const OpDef& graph_op_def, const ApiDef& api_def, // safe, fall back to the original name when THAT is safe, and only // sanitize as a last resort. const std::string safe_attr_name = - tensorflow::IsValidAttrOrArgName(api_def_attr.rename_to()) - ? api_def_attr.rename_to() - : tensorflow::IsValidAttrOrArgName(attr.name()) - ? attr.name() - : tensorflow::SanitizeToIdentifier(attr.name()); + SafeRenameTo(attr.name(), api_def_attr.rename_to()); std::string attr_name = AvoidCPPKeywords(safe_attr_name); std::string attr_comment; @@ -655,7 +657,10 @@ OpInfo::OpInfo(const OpDef& graph_op_def, const ApiDef& api_def, bool is_list = ArgIsList(arg); output_types.push_back( absl::StrCat("::tensorflow::", is_list ? "OutputList" : "Output")); - output_names.push_back(AvoidCPPKeywords(api_def_arg.rename_to())); + // See the safe_input_name comment above: the output arg's name and + // rename_to() need the same three-tier fallback as the input case. + output_names.push_back( + AvoidCPPKeywords(SafeRenameTo(arg.name(), api_def_arg.rename_to()))); is_list_output.push_back(is_list); } @@ -719,7 +724,15 @@ std::string OpInfo::GetOpAttrStruct() const { const auto entry = AttrTypeName(attr.type()); const auto attr_type_name = entry.first; const bool use_const = entry.second; - const std::string camel_case_name = ToCamelCase(api_def_attr.rename_to()); + // See the safe_input_name comment in OpInfo::OpInfo: attr.name() and + // api_def_attr.rename_to() are both unvalidated at OpDef-registration + // time, and both are spliced as a raw C++ identifier multiple times + // below (the setter name, the field access, the static defaults + // function name, and the field declaration) -- computed once here so + // every splice site below agrees on the same identifier. + const std::string safe_attr_name = + SafeRenameTo(attr.name(), api_def_attr.rename_to()); + const std::string camel_case_name = ToCamelCase(safe_attr_name); const std::string suffix = (camel_case_name == op_name || camel_case_name == "Attrs") ? "_" : ""; const std::string attr_func_def = @@ -738,7 +751,7 @@ std::string OpInfo::GetOpAttrStruct() const { absl::StrAppend(&setters, " TF_MUST_USE_RESULT Attrs ", attr_func_def, " x) {\n"); absl::StrAppend(&setters, " Attrs ret = *this;\n"); - absl::StrAppend(&setters, " ret.", api_def_attr.rename_to(), + absl::StrAppend(&setters, " ret.", safe_attr_name, "_ = x;\n"); absl::StrAppend(&setters, " return ret;\n }\n\n"); @@ -749,7 +762,7 @@ std::string OpInfo::GetOpAttrStruct() const { // Non-empty lists need static storage for their defaults. Define a // function with static local variable that stores the array. absl::StrAppend(&defaults_static_storage, " static ", attr_type_name, - " Default_", api_def_attr.rename_to(), "() {\n"); + " Default_", safe_attr_name, "() {\n"); absl::StrAppend( &defaults_static_storage, " static const ", ListElementTypeName(attr.type()), " kStorage[] = ", @@ -758,14 +771,14 @@ std::string OpInfo::GetOpAttrStruct() const { absl::StrAppend(&defaults_static_storage, " return ", attr_type_name, "(kStorage);\n }\n"); // Set the field_initializer to call the defined function. - absl::StrAppend(&field_initiliazer, "Default_", api_def_attr.rename_to(), + absl::StrAppend(&field_initiliazer, "Default_", safe_attr_name, "()"); } else { field_initiliazer = PrintAttrValue(graph_op_def.name(), api_def_attr.default_value()); } absl::StrAppend(&struct_fields, " ", attr_type_name, " ", - api_def_attr.rename_to(), "_ = ", field_initiliazer, ";\n"); + safe_attr_name, "_ = ", field_initiliazer, ";\n"); } if (struct_fields.empty()) { @@ -787,3 +800,4 @@ std::string OpInfo::GetOpAttrStruct() const { } // namespace cc_op } // namespace tensorflow + diff --git a/tensorflow/cc/framework/cc_op_gen_util.h b/tensorflow/cc/framework/cc_op_gen_util.h index de954fd3f9fc42..6e8374d8dea1ce 100644 --- a/tensorflow/cc/framework/cc_op_gen_util.h +++ b/tensorflow/cc/framework/cc_op_gen_util.h @@ -106,6 +106,16 @@ bool IsCPPKeyword(absl::string_view name); std::string AvoidCPPKeywords(absl::string_view name); +// Returns a name safe to splice as a raw C++ identifier: `rename_to` if it +// is already a safe identifier (see IsValidAttrOrArgName), else `name` if +// THAT is safe, else a sanitized fallback (see SanitizeToIdentifier). +// `rename_to` comes from ApiDef, a separate message from the OpDef arg/attr +// name `name` is drawn from; neither is validated at OpDef-registration +// time (see op_def_util.cc for why), so both must be checked here rather +// than trusted, at every place either is spliced as a raw C++ identifier +// into generated source. +std::string SafeRenameTo(absl::string_view name, absl::string_view rename_to); + void InferArgAttributes( const OpDef::ArgDef& arg, std::unordered_map* inferred_attrs); @@ -152,3 +162,4 @@ struct OpInfo { } // namespace tensorflow #endif // TENSORFLOW_CC_FRAMEWORK_CC_OP_GEN_UTIL_H_ + From 45e38a2729b26ad24ff416449a7c8c19f7a779b7 Mon Sep 17 00:00:00 2001 From: Vaggelis Date: Tue, 25 Aug 2026 09:32:46 -0400 Subject: [PATCH 11/32] Register missing AdjustContrastv2 gradient tf.image.adjust_contrast could not be differentiated: the raw AdjustContrastv2 op had no Python gradient registration, so any tape through it raised "LookupError: gradient registry has no entry for: AdjustContrastv2". The kernel computes (images - mean) * factor + mean, where mean is taken per batch and channel over the last three dimensions, interpreted as [height, width, channels]. The new registration reduces the incoming gradient over those same axes and returns sum(grad * (images - mean)) for the scalar factor. The half and float paths form the factor reduction in float32 to avoid half-precision overflow. Filed as issue 126083; adjust_hue and adjust_saturation need piecewise HSV derivations and are left for separate changes. Test Plan: Added AdjustContrastOpTestBase with gradient_checker_v2 cases for rank 3, 4 and 5 inputs so the analytical gradient is checked against finite differences of the real forward kernel. Ran image_grad_test.py AdjustContrastOpTest against a pip tf-nightly build with the patched image_grad.py overlaid: "Ran 4 tests in 2.486s / OK (skipped=1)". With pristine image_grad.py the rank 3 and rank 4 cases fail with the LookupError above. --- RELEASE.md | 7 ++++ tensorflow/python/ops/image_grad.py | 36 +++++++++++++++++++ tensorflow/python/ops/image_grad_test.py | 1 + tensorflow/python/ops/image_grad_test_base.py | 32 +++++++++++++++++ 4 files changed, 76 insertions(+) diff --git a/RELEASE.md b/RELEASE.md index b660edf249e541..c8a5d1bd0c17c2 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -34,6 +34,13 @@ In `tensorflow/c/experimental/filesystem/filesystem_interface.h`, removed `TF_Tr * Exports `__new__` in public API golden files for subclasses of `tuple` (like `tf.io.FixedLenFeature`) to fix false positives during static type checking.> * `tf.data` * Fixes a bug in `tf.data.Dataset.scan` where the shape of the state returned by `scan_func` was not strictly validated against the initial state. +* `tf.image.adjust_contrast` + + * Registers the missing Python gradient for the `AdjustContrastv2` op, + so `tf.image.adjust_contrast` can now be differentiated with + `GradientTape`. Fixes + [#126083](https://github.com/tensorflow/tensorflow/issues/126083). + * `tf.experimental.numpy` * `tf.experimental.numpy.isclose` and `tf.experimental.numpy.allclose` now diff --git a/tensorflow/python/ops/image_grad.py b/tensorflow/python/ops/image_grad.py index 39517bfb132d32..e8d9a706025ef5 100644 --- a/tensorflow/python/ops/image_grad.py +++ b/tensorflow/python/ops/image_grad.py @@ -47,6 +47,42 @@ def _ResizeNearestNeighborGrad(op: ops.Operation, grad): return [grads, None] +@ops.RegisterGradient("AdjustContrastv2") +def _AdjustContrastGrad(op: ops.Operation, grad): + """The derivatives for `tf.image.adjust_contrast`. + + The kernel computes `(images - mean) * contrast_factor + mean`, where + `mean` is taken per batch and channel over the last three dimensions, + which are interpreted as [height, width, channels]. + + Args: + op: The `AdjustContrastv2` `Operation`. + grad: The tensor representing the gradient w.r.t. the output. + + Returns: + The gradients w.r.t. the images and the contrast factor. + """ + images = op.inputs[0] + factor = op.inputs[1] + factor_t = math_ops.cast(factor, images.dtype) + rank = images.shape.rank + if rank is not None: + spatial_axes = list(range(rank - 3, rank - 1)) + else: + spatial_axes = math_ops.range(rank - 3, rank - 1) + mean = math_ops.reduce_mean(images, axis=spatial_axes, keepdims=True) + grad_images = grad * factor_t + (1.0 - factor_t) * math_ops.reduce_mean( + grad, + axis=spatial_axes, + keepdims=True, + ) + grad_factor = math_ops.reduce_sum( + math_ops.cast(grad, dtypes.float32) * + (math_ops.cast(images, dtypes.float32) - + math_ops.cast(mean, dtypes.float32))) + return grad_images, grad_factor + + @ops.RegisterGradient("ResizeBilinear") def _ResizeBilinearGrad(op: ops.Operation, grad): """The derivatives for bilinear resizing. diff --git a/tensorflow/python/ops/image_grad_test.py b/tensorflow/python/ops/image_grad_test.py index 50636b9beef435..28fd605e628455 100644 --- a/tensorflow/python/ops/image_grad_test.py +++ b/tensorflow/python/ops/image_grad_test.py @@ -23,6 +23,7 @@ ScaleAndTranslateOpTest = test_base.ScaleAndTranslateOpTestBase CropAndResizeOpTest = test_base.CropAndResizeOpTestBase RGBToHSVOpTest = test_base.RGBToHSVOpTestBase +AdjustContrastOpTest = test_base.AdjustContrastOpTestBase if __name__ == "__main__": test.main() diff --git a/tensorflow/python/ops/image_grad_test_base.py b/tensorflow/python/ops/image_grad_test_base.py index fc84f3054f9656..980b3aeae5261a 100644 --- a/tensorflow/python/ops/image_grad_test_base.py +++ b/tensorflow/python/ops/image_grad_test_base.py @@ -644,5 +644,37 @@ def f_dummy(x): self.assertAllClose(numerical_dummy, numerical, atol=1e-4) +class AdjustContrastOpTestBase(test.TestCase): + """Tests the gradient of tf.image.adjust_contrast. + + The op is affine in the images and linear in the scalar factor, so + gradient_checker_v2 validates both slots against the real forward kernel, + which is what pins down the spatial reduction axes. + """ + + def testGradRank3(self): + self._check_gradient([4, 5, 3]) + + def testGradRank4(self): + self._check_gradient([2, 4, 5, 3]) + + def testGradRank5(self): + self._check_gradient([2, 3, 4, 5, 3]) + + def _check_gradient(self, shape): + x = np.linspace(0.05, 0.95, num=int(np.prod(shape))).reshape(shape) + x = x.astype(np.float32) + factor = constant_op.constant(1.3, dtype=dtypes.float32) + + def f(image_tensor, factor_tensor): + return image_ops.adjust_contrast(image_tensor, factor_tensor) + + with self.cached_session(): + analytical, numerical = gradient_checker_v2.compute_gradient( + f, [constant_op.constant(x), factor]) + max_error = gradient_checker_v2.max_error(analytical, numerical) + self.assertLess(max_error, 1e-4) + + if __name__ == '__main__': test.main() From 0320f141c7944608a3ce6ebc6624a85d4b8bf1fb Mon Sep 17 00:00:00 2001 From: Vaggelis Date: Tue, 25 Aug 2026 12:00:49 -0400 Subject: [PATCH 12/32] Fix dynamic-rank fallback in AdjustContrastv2 gradient The unknown-rank branch subtracted 3 from the Python value None instead of the symbolic rank, which would raise TypeError whenever a graph placeholder of unknown rank reached it. Read the rank through array_ops.rank there, matching the intent of the branch. Caught in code review on pull request 126086. Test Plan: python -m py_compile tensorflow/python/ops/image_grad.py image_grad_test.py AdjustContrastOpTest against the nightly overlay: "Ran 4 tests in 1.393s / OK (skipped=1)". --- tensorflow/python/ops/image_grad.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tensorflow/python/ops/image_grad.py b/tensorflow/python/ops/image_grad.py index e8d9a706025ef5..ebf1bce0dbeb53 100644 --- a/tensorflow/python/ops/image_grad.py +++ b/tensorflow/python/ops/image_grad.py @@ -65,11 +65,12 @@ def _AdjustContrastGrad(op: ops.Operation, grad): images = op.inputs[0] factor = op.inputs[1] factor_t = math_ops.cast(factor, images.dtype) - rank = images.shape.rank - if rank is not None: - spatial_axes = list(range(rank - 3, rank - 1)) + static_rank = images.shape.rank + if static_rank is not None: + spatial_axes = list(range(static_rank - 3, static_rank - 1)) else: - spatial_axes = math_ops.range(rank - 3, rank - 1) + dynamic_rank = array_ops.rank(images) + spatial_axes = math_ops.range(dynamic_rank - 3, dynamic_rank - 1) mean = math_ops.reduce_mean(images, axis=spatial_axes, keepdims=True) grad_images = grad * factor_t + (1.0 - factor_t) * math_ops.reduce_mean( grad, From 660f795f9f0bf7c4345e4b510d955dcbc04b58ed Mon Sep 17 00:00:00 2001 From: Vaggelis Date: Wed, 26 Aug 2026 03:57:41 -0400 Subject: [PATCH 13/32] Regenerate pywrap_gradient_exclusions for AdjustContrastv2 The new AdjustContrastv2 registration changes the gradient exclusion tables that pywrap_gradient_exclusions.cc holds: the registration reads only the grad argument and none of op.inputs, so the op gains a full unused-inputs entry. Regenerated with the documented generator entry point; the output for unmodified master reproduces the committed file byte for byte, and the only delta here is the single new AdjustContrastv2 line. --- tensorflow/python/eager/pywrap_gradient_exclusions.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tensorflow/python/eager/pywrap_gradient_exclusions.cc b/tensorflow/python/eager/pywrap_gradient_exclusions.cc index f58e855263dfce..8c34921427963b 100644 --- a/tensorflow/python/eager/pywrap_gradient_exclusions.cc +++ b/tensorflow/python/eager/pywrap_gradient_exclusions.cc @@ -429,13 +429,14 @@ absl::optional> OpGradientUnusedInputIndices( absl::optional> OpGradientUnusedOutputIndices( const tensorflow::string &op_name) { - static std::array a = {{ + static std::array a = {{ {"Abs"}, {"AccumulateNV2"}, {"Acos"}, {"Add"}, {"AddN"}, {"AddV2"}, + {"AdjustContrastv2"}, {"AllToAll"}, {"Angle"}, {"ApproxTopK", 1, {0}}, From efbd98469279ba033ee881c14881ab73319152b0 Mon Sep 17 00:00:00 2001 From: nishad shabbir Date: Wed, 26 Aug 2026 15:14:50 +0530 Subject: [PATCH 14/32] Make `tf.experimental.numpy.fabs` always return a floating point result `fabs` was defined as `return abs(x)`, so it inherited `absolute`'s dtype behaviour and handed back an integer for an integer argument: >>> tnp.fabs(np.array([1, -2, 3], dtype=np.int32)).dtype tf.int32 >>> np.fabs(np.array([1, -2, 3], dtype=np.int32)).dtype dtype('float64') Returning a float is the whole difference between `fabs` and `absolute`, and the docstring `np_doc('fabs')` generates points at numpy's, which documents the result as always floating point. Pass `promote_to_float=True` to `_scalar()`, the same way `ceil`, `floor` and the other always-float unary ops in this file do. Floating point arguments are unaffected, since `_scalar()` only casts a dtype that is not already inexact, so `float16` and `float32` inputs keep their own dtype exactly as numpy does. The values are unchanged either way, only the dtype moves, so the new test checks the dtype rather than just the numbers. --- .../python/ops/numpy_ops/np_math_ops.py | 4 +++- .../python/ops/numpy_ops/np_math_ops_test.py | 24 +++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py index 8fa7ca81622510..ab8529c9406f4f 100644 --- a/tensorflow/python/ops/numpy_ops/np_math_ops.py +++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py @@ -784,7 +784,9 @@ def absolute(x): @tf_export.tf_export('experimental.numpy.fabs', v1=[]) @np_utils.np_doc('fabs') def fabs(x): - return abs(x) + # Unlike `absolute`, `fabs` always produces a floating point result, + # so an integer argument has to be promoted first. + return _scalar(math_ops.abs, x, True) @tf_export.tf_export('experimental.numpy.ceil', v1=[]) diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py index 641fdac10b585b..75d25ca30aca44 100644 --- a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py +++ b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py @@ -687,6 +687,30 @@ def testIsInfFamilyFloatInputs(self): ) + def testFabsAlwaysReturnsFloat(self): + # `fabs` differs from `absolute` in that its result is always floating + # point, so an integer argument is promoted rather than passed through. + int_args = [ + [1, -2, 3], + -5, + np.array([1, -2, 3], dtype=np.int32), + np.array([1, -2, 3], dtype=np.int64), + np.array([], dtype=np.int32), + np.array([[1, -2], [3, -4]], dtype=np.int32), + ] + for arg in int_args: + self.match( + np_math_ops.fabs(arg), np.fabs(arg), msg='fabs({})'.format(arg) + ) + + # A floating point argument keeps its own dtype. + for dtype in [np.float16, np.float32, np.float64]: + arg = np.array([1.5, -2.5], dtype=dtype) + self.match( + np_math_ops.fabs(arg), np.fabs(arg), msg='fabs({})'.format(arg) + ) + + if __name__ == '__main__': tensor.enable_tensor_equality() ops.enable_eager_execution() From 865e3983a92448587d96366767a75dc5c2f47e3f Mon Sep 17 00:00:00 2001 From: PrabinDevkota Date: Wed, 26 Aug 2026 20:44:57 +0530 Subject: [PATCH 15/32] Use a rank-deficient nonzero matrix in the XLA determinant test to avoid TPU QR NaNs. --- tensorflow/compiler/tests/matrix_determinant_op_test.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tensorflow/compiler/tests/matrix_determinant_op_test.py b/tensorflow/compiler/tests/matrix_determinant_op_test.py index f237eecb83f830..828718f55e490c 100644 --- a/tensorflow/compiler/tests/matrix_determinant_op_test.py +++ b/tensorflow/compiler/tests/matrix_determinant_op_test.py @@ -70,8 +70,9 @@ def testBasic(self): self._verifyDeterminantReal(np.array([[1., 2.], [3., 4.]])) # 2x2 with positive determinant (the motivating jit_compile example). self._verifyDeterminantReal(np.array([[4., 7.], [2., 6.]])) - # Singular. - self._verifyDeterminantReal(np.array([[0., 0.], [0., 0.]])) + # Rank-deficient but not all-zero. An all-zero matrix hits division by + # zero in Householder QR on TPU (SLogDet returns NaN instead of 0). + self._verifyDeterminantReal(np.array([[1., 2.], [2., 4.]])) # 3x3 with negative determinant. self._verifyDeterminantReal( np.array([[1., 0., 0.], [0., 1., 0.], [0., 0., -1.]])) From f7439e10069daee69b056993f58bae4198983b30 Mon Sep 17 00:00:00 2001 From: PrabinDevkota Date: Thu, 27 Aug 2026 00:18:37 +0530 Subject: [PATCH 16/32] Compare slogdet sign and log-abs only on non-singular XLA determinant cases. --- .../compiler/tests/matrix_determinant_op_test.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tensorflow/compiler/tests/matrix_determinant_op_test.py b/tensorflow/compiler/tests/matrix_determinant_op_test.py index 828718f55e490c..9b1f9afd85a859 100644 --- a/tensorflow/compiler/tests/matrix_determinant_op_test.py +++ b/tensorflow/compiler/tests/matrix_determinant_op_test.py @@ -57,6 +57,21 @@ def _verifyDeterminant(self, x, np_type): self.assertAllClose(np_recon, tf_recon, rtol=1e-3, atol=1e-3) self.assertShapeEqual(np_sign, sign) self.assertShapeEqual(np_log_abs, log_abs) + # Householder QR may return sign ±1 with a large negative log-abs-det for + # singular matrices (NumPy slogdet uses sign 0). Only compare components + # where NumPy reports a nonzero sign. + non_singular = np.abs(np_sign) > 0 + if np.any(non_singular): + self.assertAllClose( + np.where(non_singular, sign_out, 0), + np.where(non_singular, np_sign, 0), + rtol=1e-3, + atol=1e-3) + self.assertAllClose( + np.where(non_singular, log_abs_out, 0), + np.where(non_singular, np_log_abs, 0), + rtol=1e-3, + atol=1e-3) def _verifyDeterminantReal(self, x): for np_type in self.float_types & {np.float32, np.float64}: From 01e5bf173611fd3d9994d174192bf9924630cb0a Mon Sep 17 00:00:00 2001 From: Pat Notz Date: Wed, 26 Aug 2026 16:01:24 -0700 Subject: [PATCH 17/32] Add type checking and fix double-decref bugs in eager pywrap tape and accumulator APIs. This change adds `PyObject_TypeCheck` validation to several tape, variable watcher, and forward accumulator functions in `pywrap_tfe_src.cc` to prevent crashes when invalid Python objects are passed. It also fixes potential double-decref issues in `TFE_Py_VariableWatcherRemove` and `TFE_Py_ForwardAccumulatorSetRemove` by ensuring `Py_DECREF` is only called if the object was successfully erased from the tracking set. Finally, it updates the pybind11 wrappers to correctly propagate Python exceptions raised by these type checks and adds corresponding unit tests. PiperOrigin-RevId: 971552863 --- tensorflow/python/eager/pywrap_tfe_src.cc | 65 ++++++++++++++++++++-- tensorflow/python/eager/pywrap_tfe_test.py | 33 +++++++++++ tensorflow/python/tfe_wrapper.cc | 24 ++++++-- 3 files changed, 112 insertions(+), 10 deletions(-) diff --git a/tensorflow/python/eager/pywrap_tfe_src.cc b/tensorflow/python/eager/pywrap_tfe_src.cc index c1d4f1089e3f69..14c4432064faf9 100644 --- a/tensorflow/python/eager/pywrap_tfe_src.cc +++ b/tensorflow/python/eager/pywrap_tfe_src.cc @@ -1870,14 +1870,15 @@ class AccumulatorSet { return true; } - void erase(TFE_Py_ForwardAccumulator* element) { + bool erase(TFE_Py_ForwardAccumulator* element) { MapType::iterator existing = map_.find(element); if (existing == map_.end()) { - return; + return false; } ListType::iterator list_position = existing->second; map_.erase(existing); ordered_.erase(list_position); + return true; } bool empty() const { return ordered_.empty(); } @@ -2379,6 +2380,10 @@ void TFE_Py_TapeVariableAccessed(PyObject* variable) { } void TFE_Py_TapeWatchVariable(PyObject* tape, PyObject* variable) { + if (!PyObject_TypeCheck(tape, &TFE_Py_Tape_Type)) { + PyErr_SetString(PyExc_TypeError, "Expected a TFE_Py_Tape object"); + return; + } if (!CouldBackprop()) { return; } @@ -2386,6 +2391,10 @@ void TFE_Py_TapeWatchVariable(PyObject* tape, PyObject* variable) { } PyObject* TFE_Py_TapeWatchedVariables(PyObject* tape) { + if (!PyObject_TypeCheck(tape, &TFE_Py_Tape_Type)) { + PyErr_SetString(PyExc_TypeError, "Expected a TFE_Py_Tape object"); + return nullptr; + } return reinterpret_cast(tape)->tape->GetVariablesAsPyTuple(); } @@ -2410,11 +2419,23 @@ PyObject* TFE_Py_VariableWatcherNew() { } void TFE_Py_VariableWatcherRemove(PyObject* variable_watcher) { + if (!PyObject_TypeCheck(variable_watcher, &TFE_Py_VariableWatcher_Type)) { + PyErr_SetString(PyExc_TypeError, + "Expected a TFE_Py_VariableWatcher object"); + return; + } auto* stack = GetVariableWatcherSet(); - stack->erase(reinterpret_cast(variable_watcher)); + bool erased = false; + if (stack != nullptr) { + auto* vw = reinterpret_cast(variable_watcher); + erased = stack->erase(vw) > 0; + } // We kept a reference to the variable watcher in the set to ensure it // wouldn't get deleted under us; cleaning it up here. - Py_DECREF(variable_watcher); + // We only decref if the variable watcher was actually erased from the set. + if (erased) { + Py_DECREF(variable_watcher); + } } void TFE_Py_VariableWatcherVariableAccessed(PyObject* variable) { @@ -2424,6 +2445,11 @@ void TFE_Py_VariableWatcherVariableAccessed(PyObject* variable) { } PyObject* TFE_Py_VariableWatcherWatchedVariables(PyObject* variable_watcher) { + if (!PyObject_TypeCheck(variable_watcher, &TFE_Py_VariableWatcher_Type)) { + PyErr_SetString(PyExc_TypeError, + "Expected a TFE_Py_VariableWatcher object"); + return nullptr; + } return reinterpret_cast(variable_watcher) ->variable_watcher->GetVariablesAsPyTuple(); } @@ -2908,6 +2934,10 @@ PyObject* TFE_Py_TapeGradient(PyObject* tape, PyObject* target, PyObject* sources_raw, PyObject* unconnected_gradients, TF_Status* status) { + if (!PyObject_TypeCheck(tape, &TFE_Py_Tape_Type)) { + PyErr_SetString(PyExc_TypeError, "Expected a TFE_Py_Tape object"); + return nullptr; + } TFE_Py_Tape* tape_obj = reinterpret_cast(tape); if (!tape_obj->tape->IsPersistent()) { auto* tape_set = GetTapeSet(); @@ -3042,6 +3072,11 @@ PyObject* TFE_Py_ForwardAccumulatorNew(bool use_batch) { } PyObject* TFE_Py_ForwardAccumulatorSetAdd(PyObject* accumulator) { + if (!PyObject_TypeCheck(accumulator, &TFE_Py_ForwardAccumulator_Type)) { + PyErr_SetString(PyExc_TypeError, + "Expected a TFE_Py_ForwardAccumulator object"); + return nullptr; + } TFE_Py_ForwardAccumulator* c_accumulator( reinterpret_cast(accumulator)); c_accumulator->nesting_id = tape_nesting_id_counter.fetch_add(1); @@ -3058,16 +3093,29 @@ PyObject* TFE_Py_ForwardAccumulatorSetAdd(PyObject* accumulator) { } void TFE_Py_ForwardAccumulatorSetRemove(PyObject* accumulator) { + if (!PyObject_TypeCheck(accumulator, &TFE_Py_ForwardAccumulator_Type)) { + PyErr_SetString(PyExc_TypeError, + "Expected a TFE_Py_ForwardAccumulator object"); + return; + } auto* accumulator_set = GetAccumulatorSet(); + bool erased = false; if (accumulator_set != nullptr) { - accumulator_set->erase( + erased = accumulator_set->erase( reinterpret_cast(accumulator)); } - Py_DECREF(accumulator); + if (erased) { + Py_DECREF(accumulator); + } } void TFE_Py_ForwardAccumulatorWatch(PyObject* accumulator, PyObject* tensor, PyObject* tangent) { + if (!PyObject_TypeCheck(accumulator, &TFE_Py_ForwardAccumulator_Type)) { + PyErr_SetString(PyExc_TypeError, + "Expected a TFE_Py_ForwardAccumulator object"); + return; + } int64_t tensor_id = FastTensorId(tensor); reinterpret_cast(accumulator) ->accumulator->Watch(tensor_id, tangent); @@ -3077,6 +3125,11 @@ void TFE_Py_ForwardAccumulatorWatch(PyObject* accumulator, PyObject* tensor, // Returns a new reference to the JVP Tensor. PyObject* TFE_Py_ForwardAccumulatorJVP(PyObject* accumulator, PyObject* tensor) { + if (!PyObject_TypeCheck(accumulator, &TFE_Py_ForwardAccumulator_Type)) { + PyErr_SetString(PyExc_TypeError, + "Expected a TFE_Py_ForwardAccumulator object"); + return nullptr; + } PyObject* jvp = reinterpret_cast(accumulator) ->accumulator->FetchJVP(FastTensorId(tensor)); if (jvp == nullptr) { diff --git a/tensorflow/python/eager/pywrap_tfe_test.py b/tensorflow/python/eager/pywrap_tfe_test.py index 86beaef2200410..04c0460159902f 100644 --- a/tensorflow/python/eager/pywrap_tfe_test.py +++ b/tensorflow/python/eager/pywrap_tfe_test.py @@ -383,6 +383,39 @@ def testIntAttrThatDoesNotFitIn32Bits(self): shape, minval, maxval, "seed", seed) + def testTapeAndWatcherTypeCheckAndRefcount(self): + dummy_obj = 100 + with self.assertRaises(TypeError): + pywrap_tfe.TFE_Py_TapeWatchVariable(dummy_obj, None) + with self.assertRaises(TypeError): + pywrap_tfe.TFE_Py_TapeWatchedVariables(dummy_obj) + with self.assertRaises(TypeError): + pywrap_tfe.TFE_Py_VariableWatcherRemove(dummy_obj) + with self.assertRaises(TypeError): + pywrap_tfe.TFE_Py_VariableWatcherWatchedVariables(dummy_obj) + with self.assertRaises(TypeError): + pywrap_tfe.TFE_Py_ForwardAccumulatorSetAdd(dummy_obj) + with self.assertRaises(TypeError): + pywrap_tfe.TFE_Py_ForwardAccumulatorSetRemove(dummy_obj) + with self.assertRaises(TypeError): + pywrap_tfe.TFE_Py_ForwardAccumulatorWatch(dummy_obj, None, None) + with self.assertRaises(TypeError): + pywrap_tfe.TFE_Py_ForwardAccumulatorJVP(dummy_obj, None) + + # Double remove on a tape, watcher, or accumulator should not double decref + t = pywrap_tfe.TFE_Py_TapeSetNew(False, False) + pywrap_tfe.TFE_Py_TapeSetRemove(t) + pywrap_tfe.TFE_Py_TapeSetRemove(t) + + vw = pywrap_tfe.TFE_Py_VariableWatcherNew() + pywrap_tfe.TFE_Py_VariableWatcherRemove(vw) + pywrap_tfe.TFE_Py_VariableWatcherRemove(vw) + + acc = pywrap_tfe.TFE_Py_ForwardAccumulatorNew(False) + pywrap_tfe.TFE_Py_ForwardAccumulatorSetAdd(acc) + pywrap_tfe.TFE_Py_ForwardAccumulatorSetRemove(acc) + pywrap_tfe.TFE_Py_ForwardAccumulatorSetRemove(acc) + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/tfe_wrapper.cc b/tensorflow/python/tfe_wrapper.cc index c75bfc16a8f50b..990b4cf6a88f75 100644 --- a/tensorflow/python/tfe_wrapper.cc +++ b/tensorflow/python/tfe_wrapper.cc @@ -1292,10 +1292,16 @@ PYBIND11_MODULE(_pywrap_tfe, m) { return tensorflow::PyoOrThrow( TFE_Py_TapeSetNew(persistent.ptr(), watch_accessed_variables.ptr())); }); - m.def("TFE_Py_TapeSetAdd", - [](const py::handle& tape) { TFE_Py_TapeSetAdd(tape.ptr()); }); - m.def("TFE_Py_TapeSetRemove", - [](const py::handle& tape) { TFE_Py_TapeSetRemove(tape.ptr()); }); + m.def("TFE_Py_TapeSetAdd", [](const py::handle& tape) { + bool had_err = (PyErr_Occurred() != nullptr); + TFE_Py_TapeSetAdd(tape.ptr()); + if (!had_err && PyErr_Occurred()) throw py::error_already_set(); + }); + m.def("TFE_Py_TapeSetRemove", [](const py::handle& tape) { + bool had_err = (PyErr_Occurred() != nullptr); + TFE_Py_TapeSetRemove(tape.ptr()); + if (!had_err && PyErr_Occurred()) throw py::error_already_set(); + }); m.def("TFE_Py_TapeSetStopOnThread", &TFE_Py_TapeSetStopOnThread); m.def("TFE_Py_TapeSetRestartOnThread", &TFE_Py_TapeSetRestartOnThread); m.def("TFE_Py_TapeSetIsStopped", @@ -1355,11 +1361,15 @@ PYBIND11_MODULE(_pywrap_tfe, m) { }); m.def("TFE_Py_TapeWatch", [](const py::handle& tape, const py::handle& tensor) { + bool had_err = (PyErr_Occurred() != nullptr); TFE_Py_TapeWatch(tape.ptr(), tensor.ptr()); + if (!had_err && PyErr_Occurred()) throw py::error_already_set(); }); m.def("TFE_Py_TapeWatchVariable", [](const py::handle& tape, const py::handle& variable) { + bool had_err = (PyErr_Occurred() != nullptr); TFE_Py_TapeWatchVariable(tape.ptr(), variable.ptr()); + if (!had_err && PyErr_Occurred()) throw py::error_already_set(); }); m.def("TFE_Py_TapeWatchedVariables", [](const py::handle& tape) { return tensorflow::PyoOrThrow(TFE_Py_TapeWatchedVariables(tape.ptr())); @@ -1369,7 +1379,9 @@ PYBIND11_MODULE(_pywrap_tfe, m) { m.def("TFE_Py_VariableWatcherNew", []() { return tensorflow::PyoOrThrow(TFE_Py_VariableWatcherNew()); }); m.def("TFE_Py_VariableWatcherRemove", [](const py::handle& variable_watcher) { + bool had_err = (PyErr_Occurred() != nullptr); TFE_Py_VariableWatcherRemove(variable_watcher.ptr()); + if (!had_err && PyErr_Occurred()) throw py::error_already_set(); }); m.def("TFE_Py_VariableWatcherVariableAccessed", [](const py::handle& variable) { @@ -1392,14 +1404,18 @@ PYBIND11_MODULE(_pywrap_tfe, m) { }); m.def("TFE_Py_ForwardAccumulatorSetRemove", [](const py::handle& accumulator) { + bool had_err = (PyErr_Occurred() != nullptr); TFE_Py_ForwardAccumulatorSetRemove(accumulator.ptr()); + if (!had_err && PyErr_Occurred()) throw py::error_already_set(); }); m.def("TFE_Py_ForwardAccumulatorWatch", [](const py::handle& accumulator, const py::handle& tensor, const py::handle& tangent) { + bool had_err = (PyErr_Occurred() != nullptr); TFE_Py_ForwardAccumulatorWatch(accumulator.ptr(), tensor.ptr(), tangent.ptr()); + if (!had_err && PyErr_Occurred()) throw py::error_already_set(); }); m.def("TFE_Py_ForwardAccumulatorJVP", [](const py::handle& accumulator, const py::handle& tensor) { From 49a1acf1587aff6b19a3c51ab2074d2924709a3b Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 26 Aug 2026 16:10:06 -0700 Subject: [PATCH 18/32] Fix alias 'actual' attribute resolving to None in _pywrap_tensorflow. PiperOrigin-RevId: 971558059 --- tensorflow/python/BUILD | 1 - 1 file changed, 1 deletion(-) diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD index fbf0e48fe9e614..1ca50ef4d5da69 100644 --- a/tensorflow/python/BUILD +++ b/tensorflow/python/BUILD @@ -1649,7 +1649,6 @@ pywrap_library( # @unsorted-dict-items common_lib_version_scripts = { "tensorflow/tensorflow_cc": select({ - "//tensorflow:windows": None, "//tensorflow:macos": "//tensorflow:tf_exported_symbols.lds", "//conditions:default": "//tensorflow:tf_version_script.lds", }), From 547c83551297d5bf44153c1d0ded1fb623cb5d29 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 26 Aug 2026 16:15:37 -0700 Subject: [PATCH 19/32] Use two-pass variance calculation in BatchNormExpander to prevent negative variance with large input offsets. Fixes #118701 Reverts 2df58eec27abe0424437af728ea71ff8f1f5b757 PiperOrigin-RevId: 971561152 --- third_party/xla/xla/service/BUILD | 2 +- .../xla/xla/service/batchnorm_expander.cc | 36 +++++++------ third_party/xla/xla/tests/BUILD | 4 +- .../xla/xla/tests/batch_norm_training_test.cc | 54 +++---------------- 4 files changed, 32 insertions(+), 64 deletions(-) diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD index ab249ebdbd073f..25c6fdf0f4daca 100644 --- a/third_party/xla/xla/service/BUILD +++ b/third_party/xla/xla/service/BUILD @@ -2558,7 +2558,7 @@ cc_library( "//xla:xla_data_proto_cc", "//xla/hlo/ir:hlo", "//xla/hlo/pass:hlo_pass", - "//xla/tsl/platform:logging", + "//xla/tsl/platform:statusor", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/functional:function_ref", "@com_google_absl//absl/log:check", diff --git a/third_party/xla/xla/service/batchnorm_expander.cc b/third_party/xla/xla/service/batchnorm_expander.cc index 554e29df485515..bcdf46bdcb3f28 100644 --- a/third_party/xla/xla/service/batchnorm_expander.cc +++ b/third_party/xla/xla/service/batchnorm_expander.cc @@ -38,7 +38,7 @@ limitations under the License. #include "xla/literal_util.h" #include "xla/shape.h" #include "xla/shape_util.h" -#include "xla/tsl/platform/logging.h" +#include "xla/tsl/platform/statusor.h" #include "xla/util.h" #include "xla/xla_data.pb.h" @@ -217,32 +217,34 @@ absl::Status BatchNormExpanderVisitor::HandleBatchNormTraining( HloComputation* add_reduce_computation = GetOrCreateScalarAddComputation(ptype); + // X^2. + auto operand_squared = + add_binary(operand_shape, HloOpcode::kMultiply, operand, operand); // Sum[X]. auto sum = add(HloInstruction::CreateReduce(feature_shape, operand, zero, dimensions_without_feature, add_reduce_computation)); + // Sum[X^2]. + auto squared_sum = add(HloInstruction::CreateReduce( + feature_shape, operand_squared, zero, dimensions_without_feature, + add_reduce_computation)); + // E[X]. auto mean = add(Mean(elements_per_feature, sum, add)); auto mean_broadcasted = feature_broadcast(mean); - // X - E[X]. - auto operand_minus_mean = add_binary(operand_shape, HloOpcode::kSubtract, - operand, mean_broadcasted); - - // (X - E[X])^2. - auto operand_minus_mean_squared = - add_binary(operand_shape, HloOpcode::kMultiply, operand_minus_mean, - operand_minus_mean); + // E[X^2]. + auto square_mean = add(Mean(elements_per_feature, squared_sum, add)); - // Sum[(X - E[X])^2]. - auto squared_diff_sum = add(HloInstruction::CreateReduce( - feature_shape, operand_minus_mean_squared, zero, - dimensions_without_feature, add_reduce_computation)); + // E^2[X]. + auto mean_square = + add_binary(feature_shape, HloOpcode::kMultiply, mean, mean); - // Var[X] = E[(X - E[X])^2]. - auto var = add(Mean(elements_per_feature, squared_diff_sum, add)); + // Var[X]. + auto var = + add_binary(feature_shape, HloOpcode::kSubtract, square_mean, mean_square); auto var_broadcasted = feature_broadcast(var); @@ -253,6 +255,10 @@ absl::Status BatchNormExpanderVisitor::HandleBatchNormTraining( // 1 / Sqrt[Var[X] + epsilon]. auto rsqrt_var_add_epsilon = add(Rsqrt(var_add_epsilon)); + // X - E[X]. + auto operand_minus_mean = add_binary(operand_shape, HloOpcode::kSubtract, + operand, mean_broadcasted); + // (X - E[X]) / Sqrt[Var[X] + epsilon]. auto normalized = add_binary(operand_shape, HloOpcode::kMultiply, operand_minus_mean, rsqrt_var_add_epsilon); diff --git a/third_party/xla/xla/tests/BUILD b/third_party/xla/xla/tests/BUILD index f2a1919062e564..6b44cb146095dd 100644 --- a/third_party/xla/xla/tests/BUILD +++ b/third_party/xla/xla/tests/BUILD @@ -3485,13 +3485,13 @@ xla_test( name = "batch_norm_training_test", srcs = ["batch_norm_training_test.cc"], deps = [ - ":hlo_test_base", + ":hlo_pjrt_test_base", ":xla_internal_test_main", # fixdeps: keep "//xla:literal_util", "//xla/hlo/testlib:test", "//xla/tests:xla_test_backend_predicates", + "//xla/tsl/platform:statusor", "@com_google_absl//absl/status", - "@com_google_googletest//:gtest", ], ) diff --git a/third_party/xla/xla/tests/batch_norm_training_test.cc b/third_party/xla/xla/tests/batch_norm_training_test.cc index 9cc9dac24f3eec..fdc14558fd48f0 100644 --- a/third_party/xla/xla/tests/batch_norm_training_test.cc +++ b/third_party/xla/xla/tests/batch_norm_training_test.cc @@ -13,15 +13,14 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include #include #include "xla/tests/xla_test_backend_predicates.h" -#include #include "absl/status/status.h" #include "xla/hlo/testlib/test.h" #include "xla/literal_util.h" -#include "xla/tests/hlo_test_base.h" +#include "xla/tests/hlo_pjrt_test_base.h" +#include "xla/tsl/platform/statusor.h" namespace xla { namespace { @@ -41,14 +40,15 @@ ENTRY entry { class BatchNormTrainingTest : public HloTestBase {}; TEST_F(BatchNormTrainingTest, CorrectComputation) { - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kModuleStr)); + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(kModuleStr)); auto input = LiteralUtil::CreateR2({{1.0}, {2.0}}); auto scale = LiteralUtil::CreateR1({0.5}); auto offset = LiteralUtil::CreateR1({0.1}); - ASSERT_OK_AND_ASSIGN(auto result, - Execute(std::move(module), {&input, &scale, &offset})); + TF_ASSERT_OK_AND_ASSIGN( + auto result, Execute(std::move(module), {&input, &scale, &offset})); // Decompose result tuple auto result_tuple = result.DecomposeTuple(); @@ -78,50 +78,12 @@ TEST_F(BatchNormTrainingTest, CorrectComputation) { } } -TEST_F(BatchNormTrainingTest, LargeOffset) { - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kModuleStr)); - - auto input = - LiteralUtil::CreateR2({{10000.0f + 1.0f}, {10000.0f + 2.0f}}); - auto scale = LiteralUtil::CreateR1({0.5f}); - auto offset = LiteralUtil::CreateR1({0.1f}); - - ASSERT_OK_AND_ASSIGN(auto result, - Execute(std::move(module), {&input, &scale, &offset})); - - auto result_tuple = result.DecomposeTuple(); - - auto expected_output = - LiteralUtil::CreateR2({{-0.399003029f}, {0.599003f}}); - auto expected_batch_mean = LiteralUtil::CreateR1({10000.0f + 1.5f}); - auto expected_batch_var = LiteralUtil::CreateR1({0.25f}); - - const float tolerance = 1e-4f; - - for (int i = 0; i < expected_output.element_count(); ++i) { - EXPECT_FALSE(std::isnan(result_tuple[0].data()[i])); - EXPECT_NEAR(result_tuple[0].data()[i], - expected_output.data()[i], tolerance); - } - - for (int i = 0; i < expected_batch_mean.element_count(); ++i) { - EXPECT_NEAR(result_tuple[1].data()[i], - expected_batch_mean.data()[i], tolerance); - } - - for (int i = 0; i < expected_batch_var.element_count(); ++i) { - EXPECT_FALSE(std::isnan(result_tuple[2].data()[i])); - EXPECT_GE(result_tuple[2].data()[i], 0.0f); - EXPECT_NEAR(result_tuple[2].data()[i], - expected_batch_var.data()[i], tolerance); - } -} - TEST_F(BatchNormTrainingTest, ReturnsErrorWhenHloPassesDisabled) { if (test::DeviceTypeIsOneOf({test::kGpu, test::kInterpreter, test::kTpu})) { GTEST_SKIP(); } - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kModuleStr)); + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(kModuleStr)); auto status_or_result = Execute(std::move(module), {}, /*run_hlo_passes=*/false); From 57da842f37b702cdc4110af58bc4ce5c9d211f72 Mon Sep 17 00:00:00 2001 From: Toli Yevtushenko Date: Wed, 26 Aug 2026 16:49:24 -0700 Subject: [PATCH 20/32] Value-initialize StackHelper value member. PiperOrigin-RevId: 971577926 --- third_party/xla/xla/tpu/c_api_conversions.h | 6 +++--- .../xla/xla/tpu/c_api_conversions_test.cc | 21 +++++++++++++++++++ 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/third_party/xla/xla/tpu/c_api_conversions.h b/third_party/xla/xla/tpu/c_api_conversions.h index 4394030b089984..b93c4a461c7520 100644 --- a/third_party/xla/xla/tpu/c_api_conversions.h +++ b/third_party/xla/xla/tpu/c_api_conversions.h @@ -104,10 +104,10 @@ void Destroy(XLA_HloModuleConfig* c_config); // Helper for managing stack based C -> C++ conversions. template struct StackHelper { - explicit StackHelper() {} + explicit StackHelper() : value() {} template - explicit StackHelper(const CppType& t) { + explicit StackHelper(const CppType& t) : value() { ::ApiConverter::ToC(t, &value); } ~StackHelper() { ::ApiConverter::Destroy(&value); } @@ -117,7 +117,7 @@ struct StackHelper { return ::ApiConverter::FromC(&value); } - mutable CType value; + mutable CType value{}; }; } // namespace ApiConverter diff --git a/third_party/xla/xla/tpu/c_api_conversions_test.cc b/third_party/xla/xla/tpu/c_api_conversions_test.cc index db67de69760887..e2415ca48f9b87 100644 --- a/third_party/xla/xla/tpu/c_api_conversions_test.cc +++ b/third_party/xla/xla/tpu/c_api_conversions_test.cc @@ -323,6 +323,27 @@ TEST(ProtoHelper, EmptyProto) { stream_executor::tpu::SerializedProto_Free(serialized_proto); } +TEST(StackHelper, DefaultConstruct) { + { + StackHelper stack_shape; + EXPECT_EQ(stack_shape.value.dimensions.size, 0); + } + { + StackHelper stack_layout; + EXPECT_EQ(stack_layout.value.minor_to_major.size, 0); + } + { + StackHelper stack_tile; + EXPECT_EQ(stack_tile.value.dimensions.size, 0); + } +} + +TEST(StackHelper, Conversion) { + xla::Shape cpp_shape = xla::ShapeUtil::MakeShapeWithType({4, 3}); + StackHelper stack_shape(cpp_shape); + EXPECT_EQ(stack_shape.AsCpp(), cpp_shape); +} + // TODO(b/290654348): SE_DeviceAddressBase, SE_DeviceAddressAllocator, // SE_MaybeOwningDeviceAddress From ddb9ef4ce424de14df7e74db84108d01426f7a95 Mon Sep 17 00:00:00 2001 From: Parker Schuh Date: Wed, 26 Aug 2026 17:38:54 -0700 Subject: [PATCH 21/32] Remove linearization related CommonPjRtClient subclass functions that are implemented generically in the base class. PiperOrigin-RevId: 971601833 --- .../xla/xla/pjrt/common_pjrt_client.cc | 6 +- third_party/xla/xla/pjrt/cpu/cpu_client.cc | 39 ------ third_party/xla/xla/pjrt/cpu/cpu_client.h | 12 -- third_party/xla/xla/pjrt/cpu/raw_buffer.cc | 129 ------------------ third_party/xla/xla/pjrt/cpu/raw_buffer.h | 12 -- 5 files changed, 5 insertions(+), 193 deletions(-) diff --git a/third_party/xla/xla/pjrt/common_pjrt_client.cc b/third_party/xla/xla/pjrt/common_pjrt_client.cc index 63e0f87459d866..f3da0d85167f07 100644 --- a/third_party/xla/xla/pjrt/common_pjrt_client.cc +++ b/third_party/xla/xla/pjrt/common_pjrt_client.cc @@ -405,8 +405,12 @@ absl::StatusOr CommonPjRtClient::LinearizeIntoImpl( HostBufferSemantics::kImmutableOnlyDuringCall, device_shape, byte_strides.value_or(absl::Span()), raw_buffer)); + // TODO(parkers): IsCpuId because the linearization pool shares with the + // execute pool on cpu, so we have potential deadlocks. if (host_buffer_semantics == - HostBufferSemantics::kImmutableOnlyDuringCall) { + HostBufferSemantics::kImmutableOnlyDuringCall || + IsCpuId(platform_id()) || + memory_space->kind_id() == UnpinnedHostMemorySpace::kKindId) { if (!linearized.IsAvailable()) { tsl::BlockUntilReady(linearized.GetAsyncValue()); } diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.cc b/third_party/xla/xla/pjrt/cpu/cpu_client.cc index ecbc7338db887f..64d780de8066ec 100644 --- a/third_party/xla/xla/pjrt/cpu/cpu_client.cc +++ b/third_party/xla/xla/pjrt/cpu/cpu_client.cc @@ -1078,45 +1078,6 @@ bool PjRtCpuClient::BufferFromHostBufferSupportsZeroCopy( data, type, dims, byte_strides, shape); } -absl::StatusOr PjRtCpuClient::LinearizeHostBufferInto( - const void* data, PrimitiveType type, absl::Span dims, - std::optional> byte_strides, - HostBufferSemantics host_buffer_semantics, - absl::AnyInvocable on_done_with_host_buffer, - const xla::Shape& device_shape, PjRtRawBufferRef raw_buffer) { - if (device_shape.IsToken()) { - return PjRtDeviceEventRef(tsl::MakeAvailableAsyncValueRef()); - } - auto* cpp_buf = raw_buffer->down_cast(); - if (cpp_buf == nullptr) { - return absl::InvalidArgumentError("Not a CPU raw buffer"); - } - return cpp_buf->CopyFromHostBuffer( - data, type, dims, byte_strides, host_buffer_semantics, - std::move(on_done_with_host_buffer), device_shape, async_work_runner(), - raw_client()->eigen_intraop_pool(), - raw_client()->max_transpose_threads()); -} - -absl::StatusOr PjRtCpuClient::LinearizeInto( - const LiteralSlice& literal, const xla::Shape& device_shape, - HostBufferSemantics host_buffer_semantics, PjRtRawBufferRef raw_buffer) { - if (host_buffer_semantics == - PjRtClient::HostBufferSemantics::kImmutableOnlyDuringCall) { - return absl::UnimplementedError( - "ImmutableOnlyDuringCall semantics is not supported on CPU."); - } - if (device_shape.IsToken()) { - return PjRtDeviceEventRef(tsl::MakeAvailableAsyncValueRef()); - } - auto* cpp_buf = raw_buffer->down_cast(); - if (cpp_buf == nullptr) { - return absl::InvalidArgumentError("Not a CPU raw buffer"); - } - return cpp_buf->CopyFromLiteral(literal, device_shape.layout(), - async_work_runner()); -} - absl::StatusOr PjRtCpuExecutable::GetCompiledMemoryStats() const { const auto& buffer_assignment = cpu_executable_->buffer_assignment(); diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.h b/third_party/xla/xla/pjrt/cpu/cpu_client.h index 0bd6837ebf2b0a..2276bb60252b63 100644 --- a/third_party/xla/xla/pjrt/cpu/cpu_client.h +++ b/third_party/xla/xla/pjrt/cpu/cpu_client.h @@ -298,18 +298,6 @@ class PjRtCpuClient final : public CommonPjRtClientImpl { absl::StatusOr GetMemorySpaceKindForShape( const Shape& shape) const override; - absl::StatusOr LinearizeHostBufferInto( - const void* data, PrimitiveType type, absl::Span dims, - std::optional> byte_strides, - HostBufferSemantics host_buffer_semantics, - absl::AnyInvocable on_done_with_host_buffer, - const xla::Shape& device_shape, PjRtRawBufferRef raw_buffer) override; - - absl::StatusOr LinearizeInto( - const LiteralSlice& literal, const xla::Shape& device_shape, - HostBufferSemantics host_buffer_semantics, - PjRtRawBufferRef raw_buffer) override; - bool BufferFromHostBufferSupportsZeroCopy( const void* data, PrimitiveType type, absl::Span dims, std::optional> byte_strides, const Shape& shape, diff --git a/third_party/xla/xla/pjrt/cpu/raw_buffer.cc b/third_party/xla/xla/pjrt/cpu/raw_buffer.cc index b2c0b29ef5879c..113de407961ac9 100644 --- a/third_party/xla/xla/pjrt/cpu/raw_buffer.cc +++ b/third_party/xla/xla/pjrt/cpu/raw_buffer.cc @@ -228,135 +228,6 @@ CpuRawBuffer::CopyRawDeviceToHostAndReturnEvent( return PjRtDeviceEventRef(std::move(event)); } -absl::StatusOr CpuRawBuffer::CopyFromLiteral( - const LiteralSlice& literal, const xla::Layout& layout, - AsyncWorkRunner* async_work_runner) { - auto event = tsl::MakeConstructedAsyncValueRef(); - async_work_runner->Execute([literal, layout, event, buffer = buffer_]() { - CHECK(buffer.IsConcrete()); - const xla::Shape& shape = literal.shape(); - if (shape.IsToken()) { - } else if ((!shape.has_layout() && - !xla::LayoutUtil::IsMonotonicWithDim0Major(layout)) || - shape.layout() != layout) { - auto shape_copy = xla::ShapeUtil::MakeShape( - literal.shape().element_type(), literal.shape().dimensions()); - shape_copy.mutable_layout()->mutable_minor_to_major()->assign( - layout.minor_to_major().begin(), layout.minor_to_major().end()); - - xla::Literal literal_copy(shape_copy); - CHECK_OK(literal_copy.CopyFrom(literal)); - PackOrCopy(literal_copy.shape().element_type(), literal_copy, - buffer->untyped_data(), buffer->size_bytes()); - } else { - PackOrCopy(literal.shape().element_type(), literal, - buffer->untyped_data(), buffer->size_bytes()); - } - event.SetStateConcrete(); - }); - return PjRtDeviceEventRef(std::move(event)); -} - -absl::StatusOr CpuRawBuffer::CopyFromHostBuffer( - const void* data, PrimitiveType type, absl::Span dims, - std::optional> byte_strides, - PjRtClient::HostBufferSemantics host_buffer_semantics, - absl::AnyInvocable on_done_with_host_buffer, const Shape& shape, - AsyncWorkRunner* async_work_runner, tsl::thread::ThreadPool* thread_pool, - int max_transpose_threads) { - CommonPjRtClient* client = - absl::down_cast(memory_space()->client()); - tsl::AsyncValueRef device_buffer = buffer_; - bool has_default_layout = - !byte_strides || HasMajorToMinorLayout(type, dims, *byte_strides); - const int bit_width = primitive_util::BitWidth(type); - // Packed arrays are unpacked on host and packed on device. - bool is_packed = primitive_util::IsSubByteNonPredType(type); - - size_t byte_size = ShapeUtil::ByteSizeOf(shape); - if (is_packed) { - byte_size *= 8 / bit_width; - } - auto dst_data_ptr = device_buffer->untyped_data(); - if (!has_default_layout || - (shape.has_layout() && - !LayoutUtil::IsMonotonicWithDim0Major(shape.layout())) || - is_packed) { - // If the input array does not have a major-to-minor layout or device layout - // is not major-to-minor, transpose it into the device layout. Currently we - // choose to always do this synchronously. - // TODO(phawkins): consider performing the transpose asynchronously. - // TODO(phawkins): parallelize the transpose. - std::shared_ptr transpose; - { - absl::InlinedVector permutation(dims.size()); - if (shape.has_layout()) { - absl::c_reverse_copy(shape.layout().minor_to_major(), - permutation.begin()); - } else { - absl::c_iota(permutation, 0); - } - TransposePlan::Options options; - options.elem_size_in_bytes = primitive_util::ByteWidth(type); - options.dims = dims; - options.permutation = permutation; - if (is_packed) { - options.dest_bits_per_element = bit_width; - } - if (byte_strides) { - options.input_striding = TransposePlan::Striding{*byte_strides}; - } - if (thread_pool) { - options.num_threads = - std::min(thread_pool->NumThreads(), max_transpose_threads); - } - ABSL_ASSIGN_OR_RETURN(transpose, client->GetTransposePlan(options)); - } - std::optional)>> schedule_work; - if (thread_pool && max_transpose_threads > 1) { - schedule_work = [thread_pool](std::function work) { - thread_pool->Schedule(std::move(work)); - }; - } - transpose->Execute(data, dst_data_ptr, schedule_work); - if (on_done_with_host_buffer) { - std::move(on_done_with_host_buffer)(); - on_done_with_host_buffer = nullptr; - } - } else { - bool should_sync_copy = - host_buffer_semantics == - PjRtClient::HostBufferSemantics::kImmutableOnlyDuringCall || - (byte_size < kSmallDataTransferByteSize); - if (should_sync_copy) { - std::memcpy(dst_data_ptr, data, byte_size); - if (on_done_with_host_buffer) { - std::move(on_done_with_host_buffer)(); - on_done_with_host_buffer = nullptr; - } - } else { - tsl::AsyncValueRef copy_event = - tsl::MakeConstructedAsyncValueRef(); - auto result = PjRtDeviceEventRef(copy_event.CopyRef()); - async_work_runner->Execute([device_buffer, dst_data_ptr, data, byte_size, - copy_event = std::move(copy_event), - on_done_with_host_buffer = std::move( - on_done_with_host_buffer)]() mutable { - tsl::profiler::TraceMe traceme("H2D Dispatch"); - std::memcpy(dst_data_ptr, data, byte_size); - if (on_done_with_host_buffer) { - std::move(on_done_with_host_buffer)(); - on_done_with_host_buffer = nullptr; - } - // Signal copy is complete. - copy_event.SetStateConcrete(); - }); - return result; - } - } - return PjRtDeviceEventRef(tsl::MakeAvailableAsyncValueRef()); -} - absl::StatusOr CpuRawBuffer::MakeAllocationReadyEvent() { return PjRtDeviceEventRef(tsl::MakeAvailableAsyncValueRef()); } diff --git a/third_party/xla/xla/pjrt/cpu/raw_buffer.h b/third_party/xla/xla/pjrt/cpu/raw_buffer.h index a9a73a3d6f779d..845f0f4560ce45 100644 --- a/third_party/xla/xla/pjrt/cpu/raw_buffer.h +++ b/third_party/xla/xla/pjrt/cpu/raw_buffer.h @@ -125,20 +125,8 @@ class CpuRawBuffer : public CommonPjRtRawBufferImpl { void* dst, int64_t offset, int64_t transfer_size, PjRtDeviceEventRefVector dependencies) override; - absl::StatusOr CopyFromLiteral( - const LiteralSlice& literal, const xla::Layout& layout, - AsyncWorkRunner* async_work_runner); - absl::StatusOr MakeAllocationReadyEvent() override; - absl::StatusOr CopyFromHostBuffer( - const void* data, PrimitiveType type, absl::Span dims, - std::optional> byte_strides, - PjRtClient::HostBufferSemantics host_buffer_semantics, - absl::AnyInvocable on_done_with_host_buffer, - const Shape& shape, AsyncWorkRunner* async_work_runner, - tsl::thread::ThreadPool* thread_pool, int max_transpose_threads); - void CopyTo( PjRtRawBufferRef dst_raw_buffer, PjRtDeviceEventPromiseRef definition_event_promise, From 6f1bb043e298e488a34dcfa9f38271a5c7e7e32a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Longeri?= Date: Wed, 26 Aug 2026 18:04:47 -0700 Subject: [PATCH 22/32] [Mosaic] Add helper functions for MemRefSliceOp to determine striding properties of a given slice PiperOrigin-RevId: 971614915 --- .../xla/xla/mosaic/dialect/tpu/tpu_ops.cc | 70 +++++++++++++++++++ .../xla/xla/mosaic/dialect/tpu/tpu_ops.td | 13 ++++ 2 files changed, 83 insertions(+) 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 9b1cadc113604e..92dcb8c92f7d13 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc @@ -275,6 +275,76 @@ LogicalResult MemRefSliceOp::verify() { return success(); } +std::optional MemRefSliceOp::sliceStridesAcrossSourceTiles( + const int64_t source_size, const int64_t slice_size, + const int64_t source_tile_size, const int64_t result_tile_size, + Value offset) { + CHECK_EQ(source_tile_size % result_tile_size, 0); + DCHECK(offset == nullptr || isGuaranteedDivisible(offset, result_tile_size)); + const std::optional maybe_cst_offset = + offset ? getConstantIntValue(offset) : std::nullopt; + if (maybe_cst_offset && slice_size != ShapedType::kDynamic) { + // Fully static slice + return (*maybe_cst_offset + slice_size - 1) / source_tile_size != + *maybe_cst_offset / source_tile_size; + } + if (slice_size != ShapedType::kDynamic && slice_size <= result_tile_size) { + // We never stride at all + return false; + } + if (source_size != ShapedType::kDynamic && source_size <= source_tile_size) { + // Source has only one tile + return false; + } + if (slice_size != ShapedType::kDynamic && slice_size > source_tile_size) { + // Slice is too big to be contained in a single source tile + return true; + } + // TODO(apaszke,tlongeri): Should we relax the requirement for the shape to + // be divisible by the slice size? We need to consider if accessing the last + // partial tile is allowed or not. + if (slice_size != ShapedType::kDynamic && + source_tile_size % slice_size == 0 && + source_size != ShapedType::kDynamic && source_size % slice_size == 0 && + offset != nullptr && isGuaranteedDivisible(offset, slice_size)) { + // Slice is guaranteed to fit in a single source tile. + return false; + } + return std::nullopt; +} + +std::optional MemRefSliceOp::sliceStridesWithinSourceTiles( + const int64_t source_size, const int64_t slice_size, + const int64_t source_tile_size, const int64_t result_tile_size, + Value offset) { + CHECK_EQ(source_tile_size % result_tile_size, 0); + DCHECK(offset == nullptr || isGuaranteedDivisible(offset, result_tile_size)); + if (source_tile_size == result_tile_size) { + // The source tile isn't subdivided into result tiles + return false; + } + if (slice_size != ShapedType::kDynamic) { + if (slice_size <= result_tile_size) { + // We never stride at all + return false; + } + if (slice_size > 2 * result_tile_size) { + // We stride more than once. We've checked that the result tile is smaller + // than the source tile, so we must stride within source tiles at least + // once. + return true; + } + // We stride exactly once. Is it within or across source tiles? + if (offset != nullptr) { + if (const std::optional rem = + getRemainder(offset, source_tile_size)) { + return *rem != source_tile_size - result_tile_size; + } + } + } + return std::nullopt; +} + std::optional MemRefSliceOp::verifyOffsetAndSizeTileAlignment( std::array tc_target_shape, MemRefType source_type_override) { mlir::MemRefType source_ty = getMemRef().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 ca958bb36b3d2f..0c2391d5fb776b 100644 --- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td +++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td @@ -1199,6 +1199,19 @@ def TPU_MemRefSliceOp : TPU_Op<"memref_slice", [Pure, AttrSizedOperandSegments]> std::optional verifyOffsetAndSizeTileAlignment( std::array tc_target_shape, MemRefType source_type_override = nullptr); + + // Given a source shape size, a slice shape size, a source tile size, a + // result tile size and optionally a slice offset, the below functions + // returns a tri-state boolean (nullopt means "unknown") indicating whether + // the result tiles of the slice stride across or within the source tiles. + // The result tile should divide the source tile and the offset, if + // provided, should be divisible by the result tile size. + static std::optional sliceStridesAcrossSourceTiles( + int64_t source_size, int64_t slice_size, int64_t source_tile_size, + int64_t result_tile_size, Value offset); + static std::optional sliceStridesWithinSourceTiles( + int64_t source_size, int64_t slice_size, int64_t source_tile_size, + int64_t result_tile_size, Value offset); }]; let hasVerifier = 1; let hasCanonicalizer = 1; From 8c345955c05b3f5a3a05a8a6b044813f51b32339 Mon Sep 17 00:00:00 2001 From: Changming Sun Date: Wed, 26 Aug 2026 21:51:18 -0700 Subject: [PATCH 23/32] Validate int16 kernel parameters during prepare. There are some parameter type mismatch issues. In the TfLiteConvParams struct , the stride and dilation fields are 32-bit integers, but in the ConvParams struct , these fields are declared as 16-bit integers. It is a disconnection between TFLite's public C API and the practical kernel implementations. There is no need to support 32-bit large strides/dilations, but the kernel implementations still need to check the parameter ranges for security reasons. PiperOrigin-RevId: 971703093 --- tensorflow/lite/kernels/BUILD | 4 + tensorflow/lite/kernels/activations.cc | 4 + tensorflow/lite/kernels/activations_test.cc | 40 ++++++++++ tensorflow/lite/kernels/conv.cc | 5 ++ tensorflow/lite/kernels/conv_test.cc | 30 +++++++ tensorflow/lite/kernels/depthwise_conv.cc | 20 +++-- .../lite/kernels/depthwise_conv_test.cc | 24 ++++++ tensorflow/lite/kernels/gather.cc | 9 ++- tensorflow/lite/kernels/gather_test.cc | 42 ++++++++++ tensorflow/lite/kernels/padding.h | 14 ++++ .../perception/max_pool_with_argmax.cc | 2 + .../perception/max_pool_with_argmax_test.cc | 23 +++++- tensorflow/lite/kernels/pooling.cc | 1 + tensorflow/lite/kernels/pooling_test.cc | 36 +++++++++ tensorflow/lite/kernels/split.cc | 4 +- tensorflow/lite/kernels/split_test.cc | 35 +++++++- tensorflow/lite/kernels/split_v.cc | 9 ++- tensorflow/lite/kernels/split_v_test.cc | 39 ++++++++- tensorflow/lite/kernels/transpose_conv.cc | 49 ++++++++---- .../lite/kernels/transpose_conv_test.cc | 79 +++++++++++++++---- tensorflow/lite/kernels/unpack.cc | 18 +++-- tensorflow/lite/kernels/unpack_test.cc | 36 ++++++++- 22 files changed, 470 insertions(+), 53 deletions(-) diff --git a/tensorflow/lite/kernels/BUILD b/tensorflow/lite/kernels/BUILD index 4e1f229bf90c79..76158c739db190 100644 --- a/tensorflow/lite/kernels/BUILD +++ b/tensorflow/lite/kernels/BUILD @@ -2309,6 +2309,7 @@ cc_test( deps = [ ":test_main", ":test_util", + "//tensorflow/lite/c:c_api_types", "//tensorflow/lite/schema:schema_fbs", "@com_google_googletest//:gtest", "@flatbuffers", @@ -2478,6 +2479,7 @@ cc_test( deps = [ ":test_main", ":test_util", + "//tensorflow/lite/c:c_api_types", "//tensorflow/lite/schema:schema_fbs", "@com_google_googletest//:gtest", "@flatbuffers", @@ -2491,6 +2493,7 @@ cc_test( deps = [ ":test_main", ":test_util", + "//tensorflow/lite/c:c_api_types", "//tensorflow/lite/schema:schema_fbs", "@com_google_googletest//:gtest", "@flatbuffers", @@ -2771,6 +2774,7 @@ cc_test( deps = [ ":test_main", ":test_util", + "//tensorflow/lite/c:c_api_types", "//tensorflow/lite/schema:schema_fbs", "//tensorflow/lite/types:half", "@com_google_googletest//:gtest", diff --git a/tensorflow/lite/kernels/activations.cc b/tensorflow/lite/kernels/activations.cc index 616134df602bc0..20ffc244a31225 100644 --- a/tensorflow/lite/kernels/activations.cc +++ b/tensorflow/lite/kernels/activations.cc @@ -267,6 +267,10 @@ TfLiteStatus HardSwishPrepare(TfLiteContext* context, TfLiteNode* node) { HardSwishParams* params = &data->params; const TfLiteTensor* input; TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 0, &input)); + TF_LITE_ENSURE(context, input->params.zero_point >= INT16_MIN); + TF_LITE_ENSURE(context, input->params.zero_point <= INT16_MAX); + TF_LITE_ENSURE(context, output->params.zero_point >= INT16_MIN); + TF_LITE_ENSURE(context, output->params.zero_point <= INT16_MAX); params->input_zero_point = input->params.zero_point; params->output_zero_point = output->params.zero_point; const float input_scale = input->params.scale; diff --git a/tensorflow/lite/kernels/activations_test.cc b/tensorflow/lite/kernels/activations_test.cc index 7ce04226d9196c..533c2eff4dbe10 100644 --- a/tensorflow/lite/kernels/activations_test.cc +++ b/tensorflow/lite/kernels/activations_test.cc @@ -240,6 +240,24 @@ class QuantizedActivationsOpModel : public BaseActivationsOpModel { } }; +class PrepareOnlyHardSwishOpModel : public SingleOpModel { + public: + PrepareOnlyHardSwishOpModel(const TensorData& input, + const TensorData& output) { + input_ = AddInput(input); + output_ = AddOutput(output); + SetBuiltinOp(BuiltinOperator_HARD_SWISH, BuiltinOptions_NONE, 0); + BuildInterpreter({GetShape(input_)}, /*num_threads=*/1, + /*allow_fp32_relax_to_fp16=*/false, + /*apply_delegate=*/false, + /*allocate_and_delegate=*/false); + } + + private: + int input_; + int output_; +}; + const auto kTanhKernelMap = new std::map({ {"Reference", ops::builtin::Register_TANH_REF()}, {"GenericOptimized", ops::builtin::Register_TANH_GENERIC_OPT()}, @@ -555,6 +573,28 @@ TEST(QuantizedActivationsOpTest, HardSwish) { } } +TEST(QuantizedActivationsOpTest, HardSwishRejectsZeroPointOutsideInt16Range) { + PrepareOnlyHardSwishOpModel input_model( + /*input=*/{TensorType_INT8, + {1}, + 0.0f, + 0.0f, + 1.0f, + std::numeric_limits::max() + 1}, + /*output=*/{TensorType_INT8, {}, 0.0f, 0.0f, 1.0f, 0}); + EXPECT_EQ(input_model.AllocateTensors(), kTfLiteError); + + PrepareOnlyHardSwishOpModel output_model( + /*input=*/{TensorType_INT8, {1}, 0.0f, 0.0f, 1.0f, 0}, + /*output=*/{TensorType_INT8, + {}, + 0.0f, + 0.0f, + 1.0f, + std::numeric_limits::min() - 1}); + EXPECT_EQ(output_model.AllocateTensors(), kTfLiteError); +} + // See the comment in the reference implementation of quantized HardSwish: // A numerical issue significantly affecting ImageNet classification accuracy // with MobileNet v3 is only observable at the scale of HardSwish unit tests diff --git a/tensorflow/lite/kernels/conv.cc b/tensorflow/lite/kernels/conv.cc index e85ec46fb7dd75..6ffb218a269735 100644 --- a/tensorflow/lite/kernels/conv.cc +++ b/tensorflow/lite/kernels/conv.cc @@ -394,11 +394,15 @@ TfLiteStatus Prepare(KernelType kernel_type, TfLiteContext* context, // Validate stride values TF_LITE_ENSURE(context, params->stride_height > 0); + TF_LITE_ENSURE(context, params->stride_height <= INT16_MAX); TF_LITE_ENSURE(context, params->stride_width > 0); + TF_LITE_ENSURE(context, params->stride_width <= INT16_MAX); // Validate dilation values TF_LITE_ENSURE(context, params->dilation_height_factor > 0); + TF_LITE_ENSURE(context, params->dilation_height_factor <= INT16_MAX); TF_LITE_ENSURE(context, params->dilation_width_factor > 0); + TF_LITE_ENSURE(context, params->dilation_width_factor <= INT16_MAX); const TfLiteTensor* bias = nullptr; @@ -494,6 +498,7 @@ TfLiteStatus Prepare(KernelType kernel_type, TfLiteContext* context, params->dilation_height_factor, params->dilation_width_factor, input_height, input_width, filter_height, filter_width, padding, &out_height, &out_width, &data->padding)); + TF_LITE_ENSURE_STATUS(ValidatePaddingValuesForInt16(data->padding)); int output_spatial_elements = 0; TF_LITE_ENSURE_MSG(context, diff --git a/tensorflow/lite/kernels/conv_test.cc b/tensorflow/lite/kernels/conv_test.cc index b3fee38a3a7df8..15722e5c9a36e7 100644 --- a/tensorflow/lite/kernels/conv_test.cc +++ b/tensorflow/lite/kernels/conv_test.cc @@ -362,6 +362,36 @@ TEST(ConvolutionPrepareSecurityTest, RejectsPaddingOverflow) { EXPECT_EQ(m.AllocateTensors(), kTfLiteError); } +TEST(ConvolutionPrepareSecurityTest, RejectsParametersOutsideInt16Range) { + constexpr int kTooLarge = std::numeric_limits::max() + 1; + PrepareOnlyConvolutionOpModel stride_model( + ops::builtin::Register_CONVOLUTION_GENERIC_OPT(), + {TensorType_FLOAT32, {1, 1, 1, 1}}, {TensorType_FLOAT32, {1, 1, 1, 1}}, + {TensorType_FLOAT32, {}}, + /*stride_width=*/kTooLarge, /*stride_height=*/1); + EXPECT_EQ(stride_model.AllocateTensors(), kTfLiteError); + + PrepareOnlyConvolutionOpModel dilation_model( + ops::builtin::Register_CONVOLUTION_GENERIC_OPT(), + {TensorType_FLOAT32, {1, 1, 1, 1}}, {TensorType_FLOAT32, {1, 1, 1, 1}}, + {TensorType_FLOAT32, {}}, + /*stride_width=*/1, /*stride_height=*/1, Padding_VALID, + ActivationFunctionType_NONE, /*dilation_width_factor=*/kTooLarge); + EXPECT_EQ(dilation_model.AllocateTensors(), kTfLiteError); +} + +TEST(ConvolutionPrepareSecurityTest, RejectsPaddingOutsideInt16Range) { + constexpr int kFilterWidth = + 2 * (std::numeric_limits::max() + 1) + 1; + PrepareOnlyConvolutionOpModel m( + ops::builtin::Register_CONVOLUTION_GENERIC_OPT(), + {TensorType_FLOAT32, {1, 1, 1, 1}}, + {TensorType_FLOAT32, {1, 1, kFilterWidth, 1}}, {TensorType_FLOAT32, {}}, + /*stride_width=*/1, /*stride_height=*/1, Padding_SAME); + + EXPECT_EQ(m.AllocateTensors(), kTfLiteError); +} + TEST_P(ConvolutionOpTest, SimpleTestFloat32) { ConvolutionOpModel m(GetRegistration(), {TensorType_FLOAT32, {2, 2, 4, 1}}, {TensorType_FLOAT32, {3, 2, 2, 1}}, diff --git a/tensorflow/lite/kernels/depthwise_conv.cc b/tensorflow/lite/kernels/depthwise_conv.cc index bcbf9aebcd6847..3ab42e12aca1c1 100644 --- a/tensorflow/lite/kernels/depthwise_conv.cc +++ b/tensorflow/lite/kernels/depthwise_conv.cc @@ -166,6 +166,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { const int input_channels = SizeOfDimension(input, 3); const int output_channels = SizeOfDimension(filter, 3); TF_LITE_ENSURE_EQ(context, output_channels % input_channels, 0); + TF_LITE_ENSURE(context, output_channels / input_channels <= INT16_MAX); if (has_bias) { TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kBiasTensor, &bias)); @@ -199,6 +200,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { params->dilation_width_factor, height, width, filter_height, filter_width, padding, &out_height, &out_width, &data->padding)); + TF_LITE_ENSURE_STATUS(ValidatePaddingValuesForInt16(data->padding)); // Note that quantized inference requires that all tensors have their // parameters set. This is usually done during quantized training or @@ -331,7 +333,9 @@ TfLiteStatus ComputeDepthMultiplier(TfLiteContext* context, int num_input_channels = SizeOfDimension(input, 3); TF_LITE_ENSURE(context, num_input_channels != 0); TF_LITE_ENSURE_EQ(context, num_filter_channels % num_input_channels, 0); - *depth_multiplier = num_filter_channels / num_input_channels; + const int multiplier = num_filter_channels / num_input_channels; + TF_LITE_ENSURE(context, multiplier <= INT16_MAX); + *depth_multiplier = static_cast(multiplier); return kTfLiteOk; } @@ -504,8 +508,8 @@ TfLiteStatus EvalQuantizedPerChannel(TfLiteContext* context, TfLiteNode* node, } TfLiteStatus EvalQuantizedPerChannel16x8( - const TfLiteDepthwiseConvParams* params, const OpData* data, - const TfLiteTensor* input, const TfLiteTensor* filter, + TfLiteContext* context, const TfLiteDepthwiseConvParams* params, + const OpData* data, const TfLiteTensor* input, const TfLiteTensor* filter, const TfLiteTensor* bias, TfLiteTensor* output) { DepthwiseParams op_params; op_params.padding_type = PaddingType::kSame; @@ -515,7 +519,8 @@ TfLiteStatus EvalQuantizedPerChannel16x8( op_params.stride_height = params->stride_height; op_params.dilation_width_factor = params->dilation_width_factor; op_params.dilation_height_factor = params->dilation_height_factor; - op_params.depth_multiplier = params->depth_multiplier; + TF_LITE_ENSURE_STATUS(ComputeDepthMultiplier(context, input, filter, + &op_params.depth_multiplier)); op_params.weights_offset = 0; op_params.quantized_activation_min = data->output_activation_min; op_params.quantized_activation_max = data->output_activation_max; @@ -576,7 +581,8 @@ TfLiteStatus EvalHybridPerChannel(TfLiteContext* context, TfLiteNode* node, op_params.stride_height = params->stride_height; op_params.dilation_width_factor = params->dilation_width_factor; op_params.dilation_height_factor = params->dilation_height_factor; - op_params.depth_multiplier = params->depth_multiplier; + TF_LITE_ENSURE_STATUS(ComputeDepthMultiplier(context, input, filter, + &op_params.depth_multiplier)); op_params.weights_offset = 0; op_params.float_activation_min = output_activation_min; @@ -647,8 +653,8 @@ TfLiteStatus EvalImpl(TfLiteContext* context, TfLiteNode* node) { input, filter, bias, output); break; case kTfLiteInt16: - return EvalQuantizedPerChannel16x8(params, data, input, filter, bias, - output); + return EvalQuantizedPerChannel16x8(context, params, data, input, filter, + bias, output); break; default: TF_LITE_KERNEL_LOG(context, "Type %d not currently supported.", diff --git a/tensorflow/lite/kernels/depthwise_conv_test.cc b/tensorflow/lite/kernels/depthwise_conv_test.cc index 5cd798c024360a..75f08878b7a35d 100644 --- a/tensorflow/lite/kernels/depthwise_conv_test.cc +++ b/tensorflow/lite/kernels/depthwise_conv_test.cc @@ -259,6 +259,30 @@ TEST(DepthwiseConvolutionPrepareSecurityTest, EXPECT_EQ(stride_model.AllocateTensors(), kTfLiteError); } +TEST(DepthwiseConvolutionPrepareSecurityTest, RejectsPaddingOutsideInt16Range) { + constexpr int kFilterWidth = + 2 * (std::numeric_limits::max() + 1) + 1; + PrepareOnlyDepthwiseConvolutionOpModel m( + ops::builtin::Register_DEPTHWISE_CONVOLUTION_GENERIC_OPT(), + {TensorType_FLOAT32, {1, 1, 1, 1}}, + {TensorType_FLOAT32, {1, 1, kFilterWidth, 1}}, {TensorType_FLOAT32, {}}, + Padding_SAME); + + EXPECT_EQ(m.AllocateTensors(), kTfLiteError); +} + +TEST(DepthwiseConvolutionPrepareSecurityTest, + RejectsDepthMultiplierOutsideInt16Range) { + constexpr int kTooLarge = std::numeric_limits::max() + 1; + PrepareOnlyDepthwiseConvolutionOpModel m( + ops::builtin::Register_DEPTHWISE_CONVOLUTION_GENERIC_OPT(), + {TensorType_FLOAT32, {1, 1, 1, 1}}, + {TensorType_FLOAT32, {1, 1, 1, kTooLarge}}, {TensorType_FLOAT32, {}}, + Padding_VALID); + + EXPECT_EQ(m.AllocateTensors(), kTfLiteError); +} + TEST_P(DepthwiseConvolutionOpTest, ActivationReluTest) { DepthwiseConvolutionOpModel m( GetRegistration(), {TensorType_FLOAT32, {1, 3, 2, 2}}, diff --git a/tensorflow/lite/kernels/gather.cc b/tensorflow/lite/kernels/gather.cc index d076449de21390..103949da7002dc 100644 --- a/tensorflow/lite/kernels/gather.cc +++ b/tensorflow/lite/kernels/gather.cc @@ -105,12 +105,14 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { } int axis = params->axis; + TF_LITE_ENSURE(context, axis >= INT16_MIN && axis <= INT16_MAX); if (axis < 0) { axis += NumDimensions(input); } TF_LITE_ENSURE(context, 0 <= axis && axis < NumDimensions(input)); int batch_dims = params->batch_dims; + TF_LITE_ENSURE(context, batch_dims >= INT16_MIN && batch_dims <= INT16_MAX); // batch_dims should be in range: [-rank(positions), rank(positions)]. // Negative batch_dims is added with rank of positions. if (batch_dims < 0) { @@ -163,9 +165,12 @@ TfLiteStatus Gather(TfLiteContext* context, const TfLiteGatherParams& params, } TF_LITE_ENSURE(context, indices_has_only_positive_elements); + TF_LITE_ENSURE(context, params.axis >= INT16_MIN && params.axis <= INT16_MAX); + TF_LITE_ENSURE(context, params.batch_dims >= INT16_MIN && + params.batch_dims <= INT16_MAX); tflite::GatherParams op_params; - op_params.axis = params.axis; - op_params.batch_dims = params.batch_dims; + op_params.axis = static_cast(params.axis); + op_params.batch_dims = static_cast(params.batch_dims); return optimized_ops::Gather( op_params, GetTensorShape(input), GetTensorData(input), GetTensorShape(positions), GetTensorData(positions), diff --git a/tensorflow/lite/kernels/gather_test.cc b/tensorflow/lite/kernels/gather_test.cc index f5dd5bda80b3de..40148a57bfca8a 100644 --- a/tensorflow/lite/kernels/gather_test.cc +++ b/tensorflow/lite/kernels/gather_test.cc @@ -17,6 +17,7 @@ limitations under the License. #include #include #include +#include #include #include #include @@ -136,6 +137,47 @@ class GatherOpModel : public SingleOpModel { int output_; }; +class PrepareOnlyGatherOpModel : public SingleOpModel { + public: + PrepareOnlyGatherOpModel(int axis, int batch_dims) { + input_ = AddInput({TensorType_FLOAT32, {1}}); + positions_ = AddInput({TensorType_INT32, {1}}); + output_ = AddOutput(TensorType_FLOAT32); + SetBuiltinOp(BuiltinOperator_GATHER, BuiltinOptions_GatherOptions, + CreateGatherOptions(builder_, axis, batch_dims).Union()); + BuildInterpreter({GetShape(input_), GetShape(positions_)}, + /*num_threads=*/1, + /*allow_fp32_relax_to_fp16=*/false, + /*apply_delegate=*/false, + /*allocate_and_delegate=*/false); + } + + private: + int input_; + int positions_; + int output_; +}; + +TEST(GatherPrepareTest, RejectsAxisOutsideInt16Range) { + PrepareOnlyGatherOpModel too_large(std::numeric_limits::max() + 1, + /*batch_dims=*/0); + EXPECT_EQ(too_large.AllocateTensors(), kTfLiteError); + + PrepareOnlyGatherOpModel too_small(std::numeric_limits::min() - 1, + /*batch_dims=*/0); + EXPECT_EQ(too_small.AllocateTensors(), kTfLiteError); +} + +TEST(GatherPrepareTest, RejectsBatchDimsOutsideInt16Range) { + PrepareOnlyGatherOpModel too_large( + /*axis=*/0, std::numeric_limits::max() + 1); + EXPECT_EQ(too_large.AllocateTensors(), kTfLiteError); + + PrepareOnlyGatherOpModel too_small( + /*axis=*/0, std::numeric_limits::min() - 1); + EXPECT_EQ(too_small.AllocateTensors(), kTfLiteError); +} + struct GatherOpTest : public testing::TestWithParam {}; INSTANTIATE_TEST_SUITE_P(ConstantTensor, GatherOpTest, testing::Bool()); diff --git a/tensorflow/lite/kernels/padding.h b/tensorflow/lite/kernels/padding.h index e51d683b70d36c..1c63665291aa8c 100644 --- a/tensorflow/lite/kernels/padding.h +++ b/tensorflow/lite/kernels/padding.h @@ -33,6 +33,20 @@ inline TfLiteStatus CheckedNarrowPaddingValue(int64_t value, int* result) { return kTfLiteOk; } +inline TfLiteStatus ValidatePaddingValuesForInt16( + const TfLitePaddingValues& padding_values) { + const int min = std::numeric_limits::min(); + const int max = std::numeric_limits::max(); + if (padding_values.width < min || padding_values.width > max || + padding_values.height < min || padding_values.height > max || + padding_values.width_offset < min || padding_values.width_offset > max || + padding_values.height_offset < min || + padding_values.height_offset > max) { + return kTfLiteError; + } + return kTfLiteOk; +} + inline int64_t ComputeEffectiveFilterSize(int filter_size, int dilation_rate) { return (static_cast(filter_size) - 1) * dilation_rate + 1; } diff --git a/tensorflow/lite/kernels/perception/max_pool_with_argmax.cc b/tensorflow/lite/kernels/perception/max_pool_with_argmax.cc index cb0eb842000821..0db7022a490a94 100644 --- a/tensorflow/lite/kernels/perception/max_pool_with_argmax.cc +++ b/tensorflow/lite/kernels/perception/max_pool_with_argmax.cc @@ -191,6 +191,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { op_data->params.stride_height, op_data->params.stride_width, 1, 1, height, width, op_data->params.filter_height, op_data->params.filter_width, op_data->params.padding, &out_height, &out_width); + TF_LITE_ENSURE_STATUS( + ValidatePaddingValuesForInt16(op_data->params.computed.padding)); TfLiteIntArray* output_size = TfLiteIntArrayCreate(4); output_size->data[0] = batches; diff --git a/tensorflow/lite/kernels/perception/max_pool_with_argmax_test.cc b/tensorflow/lite/kernels/perception/max_pool_with_argmax_test.cc index b87bbd8be4a8f3..39fcee09192755 100644 --- a/tensorflow/lite/kernels/perception/max_pool_with_argmax_test.cc +++ b/tensorflow/lite/kernels/perception/max_pool_with_argmax_test.cc @@ -15,6 +15,7 @@ limitations under the License. #include #include +#include #include #include @@ -41,7 +42,8 @@ class MaxpoolingWithArgMaxOpModel : public SingleOpModel { int stride_width, int filter_height, int filter_width, TfLitePadding padding, const TensorData& output, - const TensorData& indices) { + const TensorData& indices, + bool allocate_and_delegate = true) { input_ = AddInput(input); output_ = AddOutput(output); indices_ = AddOutput(indices); @@ -49,7 +51,9 @@ class MaxpoolingWithArgMaxOpModel : public SingleOpModel { std::vector custom_option = CreateCustomOptions( stride_height, stride_width, filter_height, filter_width, padding); SetCustomOp("MaxPoolWithArgmax", custom_option, RegisterMaxPoolWithArgmax); - BuildInterpreter({GetShape(input_)}); + BuildInterpreter({GetShape(input_)}, /*num_threads=*/-1, + /*allow_fp32_relax_to_fp16=*/false, + /*apply_delegate=*/true, allocate_and_delegate); } void SetInput(const std::vector& data) { @@ -102,6 +106,21 @@ class MaxpoolingWithArgMaxOpModel : public SingleOpModel { } }; +TEST(MaxpoolWithArgMaxPrepareTest, RejectsPaddingOutsideInt16Range) { + constexpr int kFilterWidth = + 2 * (std::numeric_limits::max() + 1) + 1; + MaxpoolingWithArgMaxOpModel model( + /*input=*/{TensorType_FLOAT32, {1, 1, 1, 1}}, + /*stride_height=*/1, /*stride_width=*/1, + /*filter_height=*/1, /*filter_width=*/kFilterWidth, + /*padding=*/kTfLitePaddingSame, + /*output=*/{TensorType_FLOAT32, {}}, + /*indices=*/{TensorType_INT32, {}}, + /*allocate_and_delegate=*/false); + + EXPECT_EQ(model.AllocateTensors(), kTfLiteError); +} + TEST(MaxpoolWithArgMaxTest, UnsupportedInt64Test) { EXPECT_DEATH_IF_SUPPORTED(MaxpoolingWithArgMaxOpModel model( /*input=*/{TensorType_FLOAT32, {1, 2, 4, 1}}, diff --git a/tensorflow/lite/kernels/pooling.cc b/tensorflow/lite/kernels/pooling.cc index c6e8452de1e43a..82f46da100fb90 100644 --- a/tensorflow/lite/kernels/pooling.cc +++ b/tensorflow/lite/kernels/pooling.cc @@ -95,6 +95,7 @@ TfLiteStatus GenericPrepare(TfLiteContext* context, TfLiteNode* node) { params->stride_height, params->stride_width, 1, 1, height, width, params->filter_height, params->filter_width, padding, &out_height, &out_width); + TF_LITE_ENSURE_STATUS(ValidatePaddingValuesForInt16(data->padding)); if (input->type == kTfLiteUInt8 || input->type == kTfLiteInt8) { if (pool_type == kAverage || pool_type == kMax) { diff --git a/tensorflow/lite/kernels/pooling_test.cc b/tensorflow/lite/kernels/pooling_test.cc index 4634c5ce5e7adf..45b6494c72c5f1 100644 --- a/tensorflow/lite/kernels/pooling_test.cc +++ b/tensorflow/lite/kernels/pooling_test.cc @@ -15,11 +15,13 @@ limitations under the License. #include #include +#include #include #include #include #include "flatbuffers/flatbuffers.h" // from @flatbuffers +#include "tensorflow/lite/c/c_api_types.h" #include "tensorflow/lite/kernels/test_util.h" #include "tensorflow/lite/schema/schema_generated.h" @@ -119,6 +121,40 @@ class SymmetricQuantizedPoolingOpModel16 : public BasePoolingOpModel { } }; +class PrepareOnlyPoolingOpModel : public SingleOpModel { + public: + PrepareOnlyPoolingOpModel(BuiltinOperator type, const TensorData& input, + int filter_width, int filter_height, + Padding padding = Padding_VALID) { + input_ = AddInput(input); + output_ = AddOutput({input.type, {}}); + SetBuiltinOp(type, BuiltinOptions_Pool2DOptions, + CreatePool2DOptions(builder_, padding, /*stride_w=*/1, + /*stride_h=*/1, filter_width, + filter_height, ActivationFunctionType_NONE) + .Union()); + BuildInterpreter({GetShape(input_)}, /*num_threads=*/1, + /*allow_fp32_relax_to_fp16=*/false, + /*apply_delegate=*/false, + /*allocate_and_delegate=*/false); + } + + private: + int input_; + int output_; +}; + +TEST(PoolingPrepareSecurityTest, RejectsPaddingOutsideInt16Range) { + constexpr int kFilterWidth = + 2 * (std::numeric_limits::max() + 1) + 1; + PrepareOnlyPoolingOpModel m(BuiltinOperator_AVERAGE_POOL_2D, + /*input=*/{TensorType_FLOAT32, {1, 1, 1, 1}}, + kFilterWidth, + /*filter_height=*/1, Padding_SAME); + + EXPECT_EQ(m.AllocateTensors(), kTfLiteError); +} + // Replicate each entry in a vector n times along depth (innermost dimension). // The values are incremented by delta, creating ramps offset by each input // value. This is used to create simple and predicatable variation. diff --git a/tensorflow/lite/kernels/split.cc b/tensorflow/lite/kernels/split.cc index 4b9d4178fdb507..c1fea0cc766f50 100644 --- a/tensorflow/lite/kernels/split.cc +++ b/tensorflow/lite/kernels/split.cc @@ -58,6 +58,7 @@ TfLiteStatus ResizeOutputTensors(TfLiteContext* context, TfLiteNode* node, TF_LITE_ENSURE(context, axis_value >= 0); TF_LITE_ENSURE(context, axis_value < NumDimensions(input)); + TF_LITE_ENSURE(context, axis_value <= INT16_MAX); const int input_size = SizeOfDimension(input, axis_value); TF_LITE_ENSURE(context, num_splits != 0); @@ -128,6 +129,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { TF_LITE_ENSURE(context, axis_value >= 0); TF_LITE_ENSURE(context, axis_value < NumDimensions(op_context.input)); + TF_LITE_ENSURE(context, axis_value <= INT16_MAX); // TODO(b/173221795): Our usage of VectorOfTensors could be optimized by // calculating it in Prepare, unless we defer shape calculation. @@ -138,7 +140,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { VectorOfTensors all_outputs(*context, *node->outputs); \ tflite::SplitParams op_params; \ op_params.num_split = NumOutputs(node); \ - op_params.axis = axis_value; \ + op_params.axis = static_cast(axis_value); \ reference_ops::Split(op_params, GetTensorShape(op_context.input), \ GetTensorData(op_context.input), \ all_outputs.shapes(), all_outputs.data()); \ diff --git a/tensorflow/lite/kernels/split_test.cc b/tensorflow/lite/kernels/split_test.cc index b971732ded6c9d..c7c47b66ac356a 100644 --- a/tensorflow/lite/kernels/split_test.cc +++ b/tensorflow/lite/kernels/split_test.cc @@ -15,12 +15,13 @@ limitations under the License. #include #include +#include #include #include #include #include -#include "flatbuffers/flatbuffers.h" // from @flatbuffers +#include "tensorflow/lite/c/c_api_types.h" #include "tensorflow/lite/kernels/test_util.h" #include "tensorflow/lite/schema/schema_generated.h" @@ -76,6 +77,38 @@ class SplitOpModel : public SingleOpModel { std::vector outputs_; }; +class PrepareOnlySplitOpModel : public SingleOpModel { + public: + explicit PrepareOnlySplitOpModel(const std::vector& input_shape, + int axis) { + axis_ = AddConstInput(TensorType_INT32, {axis}, {1}); + input_ = AddInput({TensorType_FLOAT32, input_shape}); + AddOutput(TensorType_FLOAT32); + SetBuiltinOp(BuiltinOperator_SPLIT, BuiltinOptions_SplitOptions, + CreateSplitOptions(builder_, /*num_splits=*/1).Union()); + BuildInterpreter({{}, GetShape(input_)}, /*num_threads=*/1, + /*allow_fp32_relax_to_fp16=*/false, + /*apply_delegate=*/false, + /*allocate_and_delegate=*/false); + } + + private: + int input_; + int axis_; +}; + +TEST(SplitAxisRangeTest, RejectsAxisOutsideInt16Range) { + constexpr int kAxis = std::numeric_limits::max() + 1; + const std::vector input_shape(kAxis + 1, 1); + PrepareOnlySplitOpModel const_axis_model(input_shape, kAxis); + EXPECT_EQ(const_axis_model.AllocateTensors(), kTfLiteError); + + SplitOpModel dynamic_axis_model({TensorType_FLOAT32, input_shape}, + /*num_splits=*/1); + dynamic_axis_model.SetAxis(/*axis=*/-1); + EXPECT_EQ(dynamic_axis_model.Invoke(), kTfLiteError); +} + template void Check(TestType test_type, int axis, int num_splits, std::initializer_list input_shape, diff --git a/tensorflow/lite/kernels/split_v.cc b/tensorflow/lite/kernels/split_v.cc index a86a35f474b671..e7c980e0bd472c 100644 --- a/tensorflow/lite/kernels/split_v.cc +++ b/tensorflow/lite/kernels/split_v.cc @@ -99,6 +99,7 @@ TfLiteStatus ResizeOutputTensors(TfLiteContext* context, TfLiteNode* node, TF_LITE_ENSURE(context, axis_value >= 0); TF_LITE_ENSURE(context, axis_value < NumDimensions(input)); + TF_LITE_ENSURE(context, axis_value <= INT16_MAX); const int input_size = SizeOfDimension(input, axis_value); if (minus_one_index != -1) { @@ -180,13 +181,19 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { } int axis_value = GetTensorData(op_context.axis)[0]; + if (axis_value < 0) { + axis_value += NumDimensions(op_context.input); + } + TF_LITE_ENSURE(context, axis_value >= 0); + TF_LITE_ENSURE(context, axis_value < NumDimensions(op_context.input)); + TF_LITE_ENSURE(context, axis_value <= INT16_MAX); // Use split function to build the outputs since they share the same logic. #define TF_LITE_SPLIT_V(scalar) \ VectorOfTensors all_outputs(*context, *node->outputs); \ tflite::SplitParams op_params; \ op_params.num_split = NumOutputs(node); \ - op_params.axis = axis_value; \ + op_params.axis = static_cast(axis_value); \ reference_ops::Split(op_params, GetTensorShape(op_context.input), \ GetTensorData(op_context.input), \ all_outputs.shapes(), all_outputs.data()); diff --git a/tensorflow/lite/kernels/split_v_test.cc b/tensorflow/lite/kernels/split_v_test.cc index c856325b55eb8f..b5f8fb3ac1c3e5 100644 --- a/tensorflow/lite/kernels/split_v_test.cc +++ b/tensorflow/lite/kernels/split_v_test.cc @@ -15,11 +15,12 @@ limitations under the License. #include #include +#include #include #include #include -#include "flatbuffers/flatbuffers.h" // from @flatbuffers +#include "tensorflow/lite/c/c_api_types.h" #include "tensorflow/lite/kernels/test_util.h" #include "tensorflow/lite/schema/schema_generated.h" @@ -87,6 +88,42 @@ class SplitVOpModel : public SingleOpModel { std::vector outputs_; }; +class PrepareOnlySplitVOpModel : public SingleOpModel { + public: + explicit PrepareOnlySplitVOpModel(const std::vector& input_shape, + int axis) { + input_ = AddInput({TensorType_FLOAT32, input_shape}); + size_splits_ = AddConstInput(TensorType_INT32, {1}, {1}); + axis_ = AddConstInput(TensorType_INT32, {axis}, {1}); + AddOutput(TensorType_FLOAT32); + SetBuiltinOp(BuiltinOperator_SPLIT_V, BuiltinOptions_SplitVOptions, + CreateSplitVOptions(builder_, /*num_splits=*/1).Union()); + BuildInterpreter({GetShape(input_), {}, {}}, /*num_threads=*/1, + /*allow_fp32_relax_to_fp16=*/false, + /*apply_delegate=*/false, + /*allocate_and_delegate=*/false); + } + + private: + int input_; + int size_splits_; + int axis_; +}; + +TEST(SplitVAxisRangeTest, RejectsAxisOutsideInt16Range) { + constexpr int kAxis = std::numeric_limits::max() + 1; + const std::vector input_shape(kAxis + 1, 1); + PrepareOnlySplitVOpModel const_axis_model(input_shape, kAxis); + EXPECT_EQ(const_axis_model.AllocateTensors(), kTfLiteError); + + SplitVOpModel dynamic_axis_model({TensorType_FLOAT32, input_shape}, + {TensorType_INT32, {1}}, + /*num_splits=*/1, kAxisIsATensor, {}); + dynamic_axis_model.SetSizeSplits({1}); + dynamic_axis_model.SetAxis(/*axis=*/-1); + EXPECT_EQ(dynamic_axis_model.Invoke(), kTfLiteError); +} + template void Check(TestType test_type, int axis, std::initializer_list input_shape, std::initializer_list size_splits_shape, diff --git a/tensorflow/lite/kernels/transpose_conv.cc b/tensorflow/lite/kernels/transpose_conv.cc index 66b18585002cf6..49300e61772483 100644 --- a/tensorflow/lite/kernels/transpose_conv.cc +++ b/tensorflow/lite/kernels/transpose_conv.cc @@ -29,9 +29,7 @@ limitations under the License. // NOLINTNEXTLINE - This header file shouldn't go to the top. #include "tensorflow/lite/kernels/internal/portable_tensor_utils.h" #include "tensorflow/lite/kernels/internal/reference/integer_ops/transpose_conv.h" -#include "tensorflow/lite/kernels/internal/reference/reference_ops.h" #include "tensorflow/lite/kernels/internal/reference/transpose_conv.h" -#include "tensorflow/lite/kernels/internal/tensor.h" #include "tensorflow/lite/kernels/internal/tensor_ctypes.h" #include "tensorflow/lite/kernels/internal/types.h" #include "tensorflow/lite/kernels/kernel_util.h" @@ -291,6 +289,34 @@ TfLiteStatus ResizeAndTransposeWeights(TfLiteContext* context, return kTfLiteOk; } +TfLiteStatus ComputeAndValidatePadding(TfLiteContext* context, + const TfLiteTransposeConvParams* params, + const TfLiteTensor* weights, + const TfLiteTensor* input, + const TfLiteTensor* output, + TfLitePaddingValues* padding) { + TF_LITE_ENSURE_EQ(context, NumDimensions(output), 4); + TF_LITE_ENSURE_EQ(context, NumDimensions(weights), 4); + TF_LITE_ENSURE_EQ(context, NumDimensions(input), 4); + TF_LITE_ENSURE(context, SizeOfDimension(output, 1) > 0); + TF_LITE_ENSURE(context, SizeOfDimension(output, 2) > 0); + const int width = SizeOfDimension(output, 2); + const int height = SizeOfDimension(output, 1); + const int filter_width = SizeOfDimension(weights, 2); + const int filter_height = SizeOfDimension(weights, 1); + + int computed_input_height, computed_input_width; + TF_LITE_ENSURE_OK( + context, ComputePaddingHeightWidthChecked( + params->stride_height, params->stride_width, 1, 1, height, + width, filter_height, filter_width, params->padding, + &computed_input_height, &computed_input_width, padding)); + TF_LITE_ENSURE_STATUS(ValidatePaddingValuesForInt16(*padding)); + TF_LITE_ENSURE_EQ(context, computed_input_height, SizeOfDimension(input, 1)); + TF_LITE_ENSURE_EQ(context, computed_input_width, SizeOfDimension(input, 2)); + return kTfLiteOk; +} + template TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { OpData* data = reinterpret_cast(node->user_data); @@ -321,6 +347,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { // Tensor sanity checks TF_LITE_ENSURE_EQ(context, NumDimensions(output_shape), 1); + TF_LITE_ENSURE_EQ(context, NumElements(output_shape), 4); TF_LITE_ENSURE_EQ(context, NumDimensions(input), 4); TF_LITE_ENSURE_EQ(context, NumDimensions(weights), 4); TF_LITE_ENSURE(context, params->stride_height > 0); @@ -406,6 +433,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { } } else { TF_LITE_ENSURE_STATUS(ResizeTensor(context, output_shape, output)); + TF_LITE_ENSURE_STATUS(ComputeAndValidatePadding( + context, params, weights, input, output, &data->padding)); if (data->has_col2im) { TF_LITE_ENSURE_STATUS( ResizeCol2ImTensor(context, output_shape, weights, input, col2im)); @@ -880,20 +909,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { weights, input, col2im)); } - // Get height and width of the output image. - const int width = SizeOfDimension(output, 2); - const int height = SizeOfDimension(output, 1); - const int filter_width = SizeOfDimension(weights, 2); - const int filter_height = SizeOfDimension(weights, 1); - - int computed_input_height, computed_input_width; - TF_LITE_ENSURE_OK(context, ComputePaddingHeightWidthChecked( - params->stride_height, params->stride_width, 1, - 1, height, width, filter_height, filter_width, - params->padding, &computed_input_height, - &computed_input_width, &data->padding)); - TF_LITE_ENSURE_EQ(context, computed_input_height, SizeOfDimension(input, 1)); - TF_LITE_ENSURE_EQ(context, computed_input_width, SizeOfDimension(input, 2)); + TF_LITE_ENSURE_STATUS(ComputeAndValidatePadding( + context, params, weights, input, output, &data->padding)); // Currently support float32, uint8, int8, int16. switch (input->type) { diff --git a/tensorflow/lite/kernels/transpose_conv_test.cc b/tensorflow/lite/kernels/transpose_conv_test.cc index 507091d74dd6df..c2e0ee1fbcba07 100644 --- a/tensorflow/lite/kernels/transpose_conv_test.cc +++ b/tensorflow/lite/kernels/transpose_conv_test.cc @@ -26,7 +26,6 @@ limitations under the License. #include #include -#include "absl/memory/memory.h" #include "tensorflow/lite/core/interpreter.h" #include "tensorflow/lite/kernels/test_util.h" #include "tensorflow/lite/schema/schema_generated.h" @@ -90,19 +89,19 @@ class BaseTransposeConvOpModel : public SingleOpModel { if (test_type == TestType::kDynamic) { PopulateTensor(output_shape_, output_shape_data); - if (!std::is_same::value && - !std::is_same::value) { + if (!std::is_same_v && + !std::is_same_v) { PopulateTensor(filter_, filter_data); } } } void SetInput(std::initializer_list data) { - if (std::is_same::value) { + if (std::is_same_v) { QuantizeAndPopulate(input_, data); - } else if (std::is_same::value) { + } else if (std::is_same_v) { QuantizeAndPopulate(input_, data); - } else if (std::is_same::value) { + } else if (std::is_same_v) { QuantizeAndPopulate(input_, data); } else { PopulateTensor(input_, data); @@ -224,6 +223,19 @@ TEST(TransposeConvPrepareSecurityTest, RejectsStrideOutsideInt16Range) { EXPECT_EQ(m.AllocateTensors(), kTfLiteError); } +TEST(TransposeConvPrepareSecurityTest, RejectsPaddingOutsideInt16Range) { + constexpr int kFilterWidth = + 2 * (std::numeric_limits::max() + 1) + 1; + PrepareOnlyTransposeConvOpModel m( + ops::builtin::Register_TRANSPOSECONV_GENERIC_OPT(), {1, 1, 1, 1}, + {TensorType_FLOAT32, {1, 1, kFilterWidth, 1}}, + {TensorType_FLOAT32, {1, 1, 1, 1}}, {TensorType_FLOAT32, {}}, + Padding_SAME, /*stride_w=*/1, /*stride_h=*/1, + ActivationFunctionType_NONE); + + EXPECT_EQ(m.AllocateTensors(), kTfLiteError); +} + TEST(TransposeConvPrepareSecurityTest, RejectsMismatchedOutputChannels) { PrepareOnlyTransposeConvOpModel m( ops::builtin::Register_TRANSPOSECONV_GENERIC_OPT(), {1, 1, 1, 2}, @@ -242,8 +254,45 @@ TEST(TransposeConvPrepareSecurityTest, RejectsInconsistentSpatialShape) { {TensorType_FLOAT32, {}}, Padding_SAME, /*stride_w=*/1, /*stride_h=*/1, ActivationFunctionType_NONE); - ASSERT_EQ(m.AllocateTensors(), kTfLiteOk); - EXPECT_EQ(m.Invoke(), kTfLiteError); + EXPECT_EQ(m.AllocateTensors(), kTfLiteError); +} + +class PrepareOnlyInvalidOutputShapeTransposeConvOpModel : public SingleOpModel { + public: + explicit PrepareOnlyInvalidOutputShapeTransposeConvOpModel( + std::initializer_list output_shape_data, + std::initializer_list output_shape_dims) { + output_shape_ = + AddConstInput(TensorType_INT32, output_shape_data, output_shape_dims); + filter_ = AddInput({TensorType_FLOAT32, {1, 1, 1, 1}}); + input_ = AddInput({TensorType_FLOAT32, {1, 1, 1, 1}}); + output_ = AddOutput({TensorType_FLOAT32, {}}); + + SetBuiltinOp( + BuiltinOperator_TRANSPOSE_CONV, BuiltinOptions_TransposeConvOptions, + CreateTransposeConvOptions(builder_, Padding_SAME, /*stride_w=*/1, + /*stride_h=*/1, ActivationFunctionType_NONE) + .Union()); + resolver_ = std::make_unique( + BuiltinOperator_TRANSPOSE_CONV, + ops::builtin::Register_TRANSPOSECONV_GENERIC_OPT(), /*version=*/1); + BuildInterpreter( + {GetShape(output_shape_), GetShape(filter_), GetShape(input_)}, + /*num_threads=*/1, /*allow_fp32_relax_to_fp16=*/false, + /*apply_delegate=*/false, + /*allocate_and_delegate=*/false); + } + + private: + int output_shape_; + int filter_; + int input_; + int output_; +}; + +TEST(TransposeConvPrepareSecurityTest, RejectsInvalidOutputShapeTensorLength) { + PrepareOnlyInvalidOutputShapeTransposeConvOpModel m({1}, {1}); + EXPECT_EQ(m.AllocateTensors(), kTfLiteError); } // Test case: @@ -983,19 +1032,19 @@ class BaseTransposeConvBiasOpModel : public SingleOpModel { GetShape(input_), GetShape(bias_)}); if (test_type == TestType::kDynamic) { PopulateTensor(output_shape_, output_shape_data); - if (!std::is_same::value && - !std::is_same::value) { + if (!std::is_same_v && + !std::is_same_v) { PopulateTensor(filter_, filter_data); } } } void SetInput(std::initializer_list data) { - if (std::is_same::value) { + if (std::is_same_v) { QuantizeAndPopulate(input_, data); - } else if (std::is_same::value) { + } else if (std::is_same_v) { QuantizeAndPopulate(input_, data); - } else if (std::is_same::value) { + } else if (std::is_same_v) { QuantizeAndPopulate(input_, data); } else { PopulateTensor(input_, data); @@ -1003,9 +1052,9 @@ class BaseTransposeConvBiasOpModel : public SingleOpModel { } void SetBias(std::initializer_list bias) { - if (std::is_same::value) { + if (std::is_same_v) { QuantizeAndPopulate(bias_, bias); - } else if (std::is_same::value) { + } else if (std::is_same_v) { PerChannelQuantizeBias(bias_, bias); } else { PopulateTensor(bias_, bias); diff --git a/tensorflow/lite/kernels/unpack.cc b/tensorflow/lite/kernels/unpack.cc index cbbff6fe964fc8..7525dc7473e4e7 100644 --- a/tensorflow/lite/kernels/unpack.cc +++ b/tensorflow/lite/kernels/unpack.cc @@ -42,10 +42,12 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input)); TF_LITE_ENSURE(context, NumElements(input) > 0); int axis = data->axis; + TF_LITE_ENSURE(context, axis >= INT16_MIN && axis <= INT16_MAX); if (axis < 0) { axis += NumDimensions(input); } TF_LITE_ENSURE(context, 0 <= axis && axis < NumDimensions(input)); + TF_LITE_ENSURE(context, axis <= INT16_MAX); if (input->type != kTfLiteInt32 && input->type != kTfLiteFloat32 && input->type != kTfLiteFloat16 && input->type != kTfLiteBFloat16 && #if defined(TFLITE_ENABLE_EXTRA_REFERENCE_KERNELS) @@ -92,7 +94,7 @@ template void UnpackImpl(TfLiteContext* context, TfLiteNode* node, const TfLiteTensor* input, int output_count, int axis) { tflite::UnpackParams op_params; - op_params.axis = axis; + op_params.axis = static_cast(axis); op_params.num_split = output_count; VectorOfTensors all_outputs(*context, *node->outputs); reference_ops::Unpack(op_params, GetTensorShape(input), @@ -106,21 +108,27 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { const TfLiteTensor* input; TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input)); + int axis = data->axis; + if (axis < 0) { + axis += NumDimensions(input); + } + TF_LITE_ENSURE(context, axis >= 0 && axis < NumDimensions(input)); + TF_LITE_ENSURE(context, axis <= INT16_MAX); switch (TfLiteTypeGetSizeBits(input->type)) { case 8: { - UnpackImpl(context, node, input, data->num, data->axis); + UnpackImpl(context, node, input, data->num, axis); break; } case 16: { - UnpackImpl(context, node, input, data->num, data->axis); + UnpackImpl(context, node, input, data->num, axis); break; } case 32: { - UnpackImpl(context, node, input, data->num, data->axis); + UnpackImpl(context, node, input, data->num, axis); break; } case 64: { - UnpackImpl(context, node, input, data->num, data->axis); + UnpackImpl(context, node, input, data->num, axis); break; } default: { diff --git a/tensorflow/lite/kernels/unpack_test.cc b/tensorflow/lite/kernels/unpack_test.cc index 89b13ce48a5baa..395603409d482a 100644 --- a/tensorflow/lite/kernels/unpack_test.cc +++ b/tensorflow/lite/kernels/unpack_test.cc @@ -15,12 +15,12 @@ limitations under the License. #include #include -#include -#include +#include #include #include #include +#include "tensorflow/lite/c/c_api_types.h" #include "tensorflow/lite/kernels/test_util.h" #include "tensorflow/lite/schema/schema_generated.h" #include "tensorflow/lite/types/half.h" @@ -77,6 +77,38 @@ class UnpackOpModel : public SingleOpModel { std::vector outputs_; }; +class PrepareOnlyUnpackOpModel : public SingleOpModel { + public: + explicit PrepareOnlyUnpackOpModel(const TensorData& input, int num_outputs, + int axis) { + input_ = AddInput(input); + for (int i = 0; i < num_outputs; ++i) { + AddOutput(input.type); + } + SetBuiltinOp(BuiltinOperator_UNPACK, BuiltinOptions_UnpackOptions, + CreateUnpackOptions(builder_, num_outputs, axis).Union()); + BuildInterpreter({GetShape(input_)}, /*num_threads=*/1, + /*allow_fp32_relax_to_fp16=*/false, + /*apply_delegate=*/false, + /*allocate_and_delegate=*/false); + } + + private: + int input_; +}; + +TEST(UnpackPrepareTest, RejectsAxisOutsideInt16Range) { + PrepareOnlyUnpackOpModel too_large( + /*input=*/{TensorType_FLOAT32, {1}}, /*num_outputs=*/1, + std::numeric_limits::max() + 1); + EXPECT_EQ(too_large.AllocateTensors(), kTfLiteError); + + PrepareOnlyUnpackOpModel too_small( + /*input=*/{TensorType_FLOAT32, {1}}, /*num_outputs=*/1, + std::numeric_limits::min() - 1); + EXPECT_EQ(too_small.AllocateTensors(), kTfLiteError); +} + template void Check(int axis, const std::initializer_list& input_shape, const std::initializer_list& input_data, From acc5b946f5d753ceef8995e12eb574a295e5e6f2 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 26 Aug 2026 21:59:34 -0700 Subject: [PATCH 24/32] Fixing typo and include warnings in memory space assignment. PiperOrigin-RevId: 971707065 --- .../service/memory_space_assignment/memory_space_assignment.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.h b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.h index 716f54fc01ec6b..f0d5c78703e8e7 100644 --- a/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.h +++ b/third_party/xla/xla/service/memory_space_assignment/memory_space_assignment.h @@ -166,7 +166,7 @@ Useful logging and error messages prefetch a buffer to alternate memory, according to some heuristic and not based on limited copy resource. * If the CostAnalysisPrefetchIntervalPicker is used, which is the default, - live range too long is governed by the picker's + live range too short is governed by the picker's min_overlap_to_async_copy_ratio argument. - "Finding allocation for": Magical logging phrase indicating the point in @@ -206,13 +206,13 @@ Useful logging and error messages #include "xla/hlo/analysis/hlo_dataflow_analysis.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/utils/hlo_live_range.h" +#include "xla/layout.h" #include "xla/service/heap_simulator/heap_simulator.h" #include "xla/service/hlo.pb.h" #include "xla/service/hlo_value.h" #include "xla/service/memory_space_assignment/allocation.h" #include "xla/service/memory_space_assignment/memory_space_assignment.pb.h" #include "xla/service/memory_space_assignment/options.h" -#include "xla/shape.h" #include "xla/util.h" namespace xla { From 20e9d377aab584e2599197e46f50b14f6a8385c8 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 26 Aug 2026 22:07:48 -0700 Subject: [PATCH 25/32] Set `element_size_in_bits` for sub-byte types in `CpuTopologyDescription::GetDefaultLayout` PiperOrigin-RevId: 971711429 --- third_party/xla/xla/pjrt/plugin/xla_cpu/BUILD | 1 + .../xla_cpu/cpu_topology_description.cc | 10 +++++++- .../xla_cpu/cpu_topology_description_test.cc | 23 +++++++++++++++++++ 3 files changed, 33 insertions(+), 1 deletion(-) diff --git a/third_party/xla/xla/pjrt/plugin/xla_cpu/BUILD b/third_party/xla/xla/pjrt/plugin/xla_cpu/BUILD index f8d22af61a690c..843696e48bef8e 100644 --- a/third_party/xla/xla/pjrt/plugin/xla_cpu/BUILD +++ b/third_party/xla/xla/pjrt/plugin/xla_cpu/BUILD @@ -145,6 +145,7 @@ xla_cc_test( deps = [ ":cpu_topology", ":cpu_topology_description", + "//xla:shape_util", "//xla/backends/cpu:target_machine_options", "//xla/pjrt:host_memory_spaces", "//xla/pjrt:pjrt_common", diff --git a/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_topology_description.cc b/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_topology_description.cc index 7a99f291db2750..088fe3053b9c08 100644 --- a/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_topology_description.cc +++ b/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_topology_description.cc @@ -77,7 +77,15 @@ absl::StatusOr CpuTopologyDescription::GetDefaultLayout( PrimitiveType_Name(element_type)); } Shape shape = ShapeUtil::MakeShape(element_type, dims); - return LayoutUtil::GetWithDefaultLayout(shape).layout(); + Layout layout = LayoutUtil::GetWithDefaultLayout(shape).layout(); + // `GetWithDefaultLayout` returns a padded layout for sub-byte types since the + // notion of "default" is context dependent and in this case means the default + // for literals for historical reasons. Because of this, we need to manually + // populate the `element_size_in_bits` for sub-byte types here. + if (primitive_util::IsSubByteNonPredType(element_type)) { + layout.set_element_size_in_bits(primitive_util::BitWidth(element_type)); + } + return layout; } absl::StatusOr CpuTopologyDescription::GetMemorySpaceKindForShape( diff --git a/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_topology_description_test.cc b/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_topology_description_test.cc index cb91e98cfca424..8d4512566a3330 100644 --- a/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_topology_description_test.cc +++ b/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_topology_description_test.cc @@ -23,11 +23,13 @@ limitations under the License. #include #include #include "xla/backends/cpu/target_machine_options.h" +#include "xla/layout.h" #include "xla/pjrt/host_memory_spaces.h" #include "xla/pjrt/pjrt_common.h" #include "xla/pjrt/pjrt_compiler.h" #include "xla/pjrt/pjrt_device_dimensions.h" #include "xla/pjrt/plugin/xla_cpu/cpu_topology.h" +#include "xla/shape_util.h" #include "xla/tsl/platform/statusor.h" namespace xla { @@ -192,5 +194,26 @@ TEST(CpuTopologyDescriptionTest, ProcessIdAndIndexOnProcess) { .ok()); } +TEST(CpuTopologyDescriptionTest, GetDefaultLayout) { + std::vector cpu_devices = {{0, 0}}; + xla::cpu::TargetMachineOptions target_machine_options( + /*triple=*/"triple", /*cpu=*/"cpu", /*features=*/""); + CpuTopologyDescription topology( + xla::CpuId(), "cpu", "1.0", + CpuTopology(cpu_devices, target_machine_options)); + + for (PrimitiveType element_type : {S4, S32}) { + const Shape shape = ShapeUtil::MakeShape(element_type, {2, 3}); + ASSERT_OK_AND_ASSIGN( + Layout default_layout, + topology.GetDefaultLayout(shape.element_type(), shape.dimensions())); + ASSERT_OK_AND_ASSIGN( + Shape canonical_shape, + topology.MakeCanonicalShapeForMemorySpace(CpuDeviceMemorySpace::kKindId, + shape, /*layout=*/nullptr)); + EXPECT_EQ(default_layout, canonical_shape.layout()); + } +} + } // namespace } // namespace xla From 0615298532f160f4f35262e9e8125cc5fd0b7242 Mon Sep 17 00:00:00 2001 From: Dmitri Latushko Date: Wed, 26 Aug 2026 22:30:34 -0700 Subject: [PATCH 26/32] Reverts 1e15c4db3f2214967aebebb69829a6c3543f6052 PiperOrigin-RevId: 971722826 --- tensorflow/compiler/jit/BUILD | 3 - ...le_functional_ops_lowering_for_xla_pass.cc | 76 ------- ...ble_functional_ops_lowering_for_xla_pass.h | 41 ---- ...nctional_ops_lowering_for_xla_pass_test.cc | 193 ------------------ .../jit/jit_compilation_pass_registration.cc | 4 - tensorflow/compiler/tests/jit_test.py | 28 --- 6 files changed, 345 deletions(-) delete mode 100644 tensorflow/compiler/jit/disable_functional_ops_lowering_for_xla_pass.cc delete mode 100644 tensorflow/compiler/jit/disable_functional_ops_lowering_for_xla_pass.h delete mode 100644 tensorflow/compiler/jit/disable_functional_ops_lowering_for_xla_pass_test.cc diff --git a/tensorflow/compiler/jit/BUILD b/tensorflow/compiler/jit/BUILD index 27dbdb1f739634..dba908f9d31d57 100644 --- a/tensorflow/compiler/jit/BUILD +++ b/tensorflow/compiler/jit/BUILD @@ -1139,7 +1139,6 @@ cc_library( "cluster_scoping_pass.cc", "deadness_analysis.cc", "deadness_analysis_internal.h", - "disable_functional_ops_lowering_for_xla_pass.cc", "encapsulate_subgraphs_pass.cc", "encapsulate_xla_computations_pass.cc", "extract_outside_compilation_pass.cc", @@ -1155,7 +1154,6 @@ cc_library( "clone_constants_for_better_clustering.h", "cluster_scoping_pass.h", "deadness_analysis.h", - "disable_functional_ops_lowering_for_xla_pass.h", "encapsulate_subgraphs_pass.h", "encapsulate_xla_computations_pass.h", "extract_outside_compilation_pass.h", @@ -1332,7 +1330,6 @@ tf_cc_test( "build_xla_ops_pass_test.cc", "clone_constants_for_better_clustering_test.cc", "cluster_scoping_pass_test.cc", - "disable_functional_ops_lowering_for_xla_pass_test.cc", "encapsulate_subgraphs_pass_test.cc", "encapsulate_xla_computations_pass_test.cc", "extract_outside_compilation_pass_test.cc", diff --git a/tensorflow/compiler/jit/disable_functional_ops_lowering_for_xla_pass.cc b/tensorflow/compiler/jit/disable_functional_ops_lowering_for_xla_pass.cc deleted file mode 100644 index e41d765276bc0b..00000000000000 --- a/tensorflow/compiler/jit/disable_functional_ops_lowering_for_xla_pass.cc +++ /dev/null @@ -1,76 +0,0 @@ -/* Copyright 2026 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/compiler/jit/disable_functional_ops_lowering_for_xla_pass.h" - -#include "tensorflow/compiler/jit/flags.h" -#include "tensorflow/compiler/jit/xla_cluster_util.h" -#include "tensorflow/core/framework/node_def_util.h" -#include "tensorflow/core/graph/graph.h" -#include "tensorflow/core/protobuf/config.pb.h" - -namespace tensorflow { - -namespace { - -// Returns the effective global JIT level for `options`. Unlike -// `GetGlobalJitLevelForGraph`, this also works when `options.session_options` -// is null, which can happen when auto-clustering is enabled purely via the -// TF_XLA_FLAGS=--tf_xla_auto_jit=N environment variable and the execution -// path never populates a SessionOptions. In that case we combine just the -// flag-provided auto-jit level with whether the graph is a single-GPU graph, -// mirroring what `GetGlobalJitLevelForGraph` does for a DEFAULT ConfigProto -// setting. -OptimizerOptions::GlobalJitLevel GetEffectiveGlobalJitLevel( - const GraphOptimizationPassOptions& options, const Graph& graph) { - if (options.session_options != nullptr) { - return GetGlobalJitLevelForGraph(options); - } - - const XlaAutoJitFlag& auto_jit_flag = - GetMarkForCompilationPassFlags()->xla_auto_jit_flag; - auto level_or_off = [](int32_t level) { - return level == OptimizerOptions::DEFAULT - ? OptimizerOptions::OFF - : static_cast(level); - }; - return IsSingleGpuGraph(graph) - ? level_or_off(auto_jit_flag.optimization_level_single_gpu) - : level_or_off(auto_jit_flag.optimization_level_general); -} - -} // namespace - -absl::Status DisableFunctionalOpsLoweringForXlaPass::Run( - const GraphOptimizationPassOptions& options) { - if (options.graph == nullptr || options.graph->get() == nullptr) { - return absl::OkStatus(); - } - Graph* graph = options.graph->get(); - if (GetEffectiveGlobalJitLevel(options, *graph) < OptimizerOptions::ON_1) { - return absl::OkStatus(); - } - for (Node* n : graph->op_nodes()) { - if (!n->IsIfNode() && !n->IsCaseNode() && !n->IsWhileNode()) continue; - bool lower = false; - if (TryGetNodeAttr(n->attrs(), "_lower_using_switch_merge", &lower) && - lower) { - n->ClearAttr("_lower_using_switch_merge"); - } - } - return absl::OkStatus(); -} - -} // namespace tensorflow diff --git a/tensorflow/compiler/jit/disable_functional_ops_lowering_for_xla_pass.h b/tensorflow/compiler/jit/disable_functional_ops_lowering_for_xla_pass.h deleted file mode 100644 index cee7eb28187c11..00000000000000 --- a/tensorflow/compiler/jit/disable_functional_ops_lowering_for_xla_pass.h +++ /dev/null @@ -1,41 +0,0 @@ -/* Copyright 2026 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef TENSORFLOW_COMPILER_JIT_DISABLE_FUNCTIONAL_OPS_LOWERING_FOR_XLA_PASS_H_ -#define TENSORFLOW_COMPILER_JIT_DISABLE_FUNCTIONAL_OPS_LOWERING_FOR_XLA_PASS_H_ - -#include "tensorflow/core/common_runtime/optimization_registry.h" - -namespace tensorflow { - -// Auto-clustering lowers functional control flow ops (While/If/Case) into -// Switch/Merge before MarkForCompilationPass runs. That can cause only -// fragments of a loop/branch body to end up in a cluster instead of the whole -// op, which silently corrupts results for some patterns. jit_compile=True -// never lowers these ops in the first place, so it doesn't hit this. -// -// This pass clears the `_lower_using_switch_merge` attribute on functional -// control flow nodes whenever global JIT is enabled, so auto-clustering skips -// the lowering the same way jit_compile=True already does. -class DisableFunctionalOpsLoweringForXlaPass : public GraphOptimizationPass { - public: - DisableFunctionalOpsLoweringForXlaPass() = default; - - absl::Status Run(const GraphOptimizationPassOptions& options) override; -}; - -} // namespace tensorflow - -#endif // TENSORFLOW_COMPILER_JIT_DISABLE_FUNCTIONAL_OPS_LOWERING_FOR_XLA_PASS_H_ diff --git a/tensorflow/compiler/jit/disable_functional_ops_lowering_for_xla_pass_test.cc b/tensorflow/compiler/jit/disable_functional_ops_lowering_for_xla_pass_test.cc deleted file mode 100644 index a0adc449288967..00000000000000 --- a/tensorflow/compiler/jit/disable_functional_ops_lowering_for_xla_pass_test.cc +++ /dev/null @@ -1,193 +0,0 @@ -/* Copyright 2026 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/compiler/jit/disable_functional_ops_lowering_for_xla_pass.h" - -#include -#include - -#include "tensorflow/compiler/jit/flags.h" -#include "tensorflow/core/framework/attr_value.pb.h" -#include "tensorflow/core/framework/node_def_builder.h" -#include "tensorflow/core/framework/node_def_util.h" -#include "tensorflow/core/framework/op.h" -#include "tensorflow/core/framework/types.h" -#include "tensorflow/core/graph/graph.h" -#include "tensorflow/core/graph/node_builder.h" -#include "tensorflow/core/lib/core/status_test_util.h" -#include "tensorflow/core/platform/test.h" -#include "tensorflow/core/protobuf/config.pb.h" -#include "tensorflow/core/public/session_options.h" - -namespace tensorflow { -namespace { - -absl::Status AddWhileNode(Graph* graph, absl::string_view name, Node** node) { - return NodeBuilder(name, "While") - .Input(std::vector{}) - .Attr("_lower_using_switch_merge", true) - .Attr("cond", NameAttrList()) - .Attr("body", NameAttrList()) - .Attr("T", DataTypeVector{}) - .Finalize(graph, node); -} - -absl::Status AddPlaceholder(Graph* graph, absl::string_view name, - DataType dtype, Node** node) { - return NodeBuilder(name, "Placeholder") - .Attr("dtype", dtype) - .Finalize(graph, node); -} - -absl::Status AddIfNode(Graph* graph, absl::string_view name, Node* cond, - Node** node) { - return NodeBuilder(name, "If") - .Input(cond, 0) - .Input(std::vector{}) - .Attr("_lower_using_switch_merge", true) - .Attr("Tcond", DT_BOOL) - .Attr("Tin", DataTypeVector{}) - .Attr("Tout", DataTypeVector{}) - .Attr("then_branch", NameAttrList()) - .Attr("else_branch", NameAttrList()) - .Finalize(graph, node); -} - -absl::Status AddCaseNode(Graph* graph, absl::string_view name, - Node* branch_index, Node** node) { - return NodeBuilder(name, "Case") - .Input(branch_index, 0) - .Input(std::vector{}) - .Attr("_lower_using_switch_merge", true) - .Attr("Tin", DataTypeVector{}) - .Attr("Tout", DataTypeVector{}) - .Attr("branches", std::vector{NameAttrList()}) - .Finalize(graph, node); -} - -bool HasLowerUsingSwitchMergeAttr(const Node* n) { - bool lower = false; - return TryGetNodeAttr(n->attrs(), "_lower_using_switch_merge", &lower) && - lower; -} - -absl::Status RunPass(std::unique_ptr* graph, - SessionOptions* session_options) { - GraphOptimizationPassOptions options; - options.graph = graph; - options.session_options = session_options; - DisableFunctionalOpsLoweringForXlaPass pass; - return pass.Run(options); -} - -TEST(DisableFunctionalOpsLoweringForXlaPassTest, - ClearsAttrOnFunctionalNodesWhenGlobalJitEnabled) { - auto graph = std::make_unique(OpRegistry::Global()); - Node* while_node; - TF_ASSERT_OK(AddWhileNode(graph.get(), "my_while", &while_node)); - Node* cond; - TF_ASSERT_OK(AddPlaceholder(graph.get(), "cond", DT_BOOL, &cond)); - Node* if_node; - TF_ASSERT_OK(AddIfNode(graph.get(), "my_if", cond, &if_node)); - Node* branch_index; - TF_ASSERT_OK( - AddPlaceholder(graph.get(), "branch_index", DT_INT32, &branch_index)); - Node* case_node; - TF_ASSERT_OK(AddCaseNode(graph.get(), "my_case", branch_index, &case_node)); - - SessionOptions session_options; - session_options.config.mutable_graph_options() - ->mutable_optimizer_options() - ->set_global_jit_level(OptimizerOptions::ON_1); - - TF_ASSERT_OK(RunPass(&graph, &session_options)); - - EXPECT_FALSE(HasLowerUsingSwitchMergeAttr(while_node)); - EXPECT_FALSE(HasLowerUsingSwitchMergeAttr(if_node)); - EXPECT_FALSE(HasLowerUsingSwitchMergeAttr(case_node)); -} - -TEST(DisableFunctionalOpsLoweringForXlaPassTest, - PreservesAttrWhenGlobalJitDisabled) { - auto graph = std::make_unique(OpRegistry::Global()); - Node* while_node; - TF_ASSERT_OK(AddWhileNode(graph.get(), "my_while", &while_node)); - - SessionOptions session_options; - session_options.config.mutable_graph_options() - ->mutable_optimizer_options() - ->set_global_jit_level(OptimizerOptions::OFF); - - TF_ASSERT_OK(RunPass(&graph, &session_options)); - - EXPECT_TRUE(HasLowerUsingSwitchMergeAttr(while_node)); -} - -TEST(DisableFunctionalOpsLoweringForXlaPassTest, NoOpWhenGraphIsMissing) { - SessionOptions session_options; - session_options.config.mutable_graph_options() - ->mutable_optimizer_options() - ->set_global_jit_level(OptimizerOptions::ON_2); - - GraphOptimizationPassOptions options; - options.session_options = &session_options; - DisableFunctionalOpsLoweringForXlaPass pass; - TF_ASSERT_OK(pass.Run(options)); -} - -class DisableFunctionalOpsLoweringForXlaPassFlagsTest : public ::testing::Test { - protected: - void SetUp() override { - flags_ = GetMarkForCompilationPassFlags(); - original_ = flags_->xla_auto_jit_flag; - } - void TearDown() override { flags_->xla_auto_jit_flag = original_; } - - MarkForCompilationPassFlags* flags_; - XlaAutoJitFlag original_; -}; - -TEST_F(DisableFunctionalOpsLoweringForXlaPassFlagsTest, - FallsBackToFlagsWhenSessionOptionsIsNull) { - auto graph = std::make_unique(OpRegistry::Global()); - Node* while_node; - TF_ASSERT_OK(AddWhileNode(graph.get(), "my_while", &while_node)); - - flags_->xla_auto_jit_flag.optimization_level_single_gpu = - OptimizerOptions::ON_2; - flags_->xla_auto_jit_flag.optimization_level_general = OptimizerOptions::ON_2; - - TF_ASSERT_OK(RunPass(&graph, /*session_options=*/nullptr)); - - EXPECT_FALSE(HasLowerUsingSwitchMergeAttr(while_node)); -} - -TEST_F(DisableFunctionalOpsLoweringForXlaPassFlagsTest, - PreservesAttrWhenSessionOptionsIsNullAndFlagsAreOff) { - auto graph = std::make_unique(OpRegistry::Global()); - Node* while_node; - TF_ASSERT_OK(AddWhileNode(graph.get(), "my_while", &while_node)); - - flags_->xla_auto_jit_flag.optimization_level_single_gpu = - OptimizerOptions::OFF; - flags_->xla_auto_jit_flag.optimization_level_general = OptimizerOptions::OFF; - - TF_ASSERT_OK(RunPass(&graph, /*session_options=*/nullptr)); - - EXPECT_TRUE(HasLowerUsingSwitchMergeAttr(while_node)); -} - -} // namespace -} // namespace tensorflow diff --git a/tensorflow/compiler/jit/jit_compilation_pass_registration.cc b/tensorflow/compiler/jit/jit_compilation_pass_registration.cc index 2abdc413e8725e..a4a9c7d7e11863 100644 --- a/tensorflow/compiler/jit/jit_compilation_pass_registration.cc +++ b/tensorflow/compiler/jit/jit_compilation_pass_registration.cc @@ -16,7 +16,6 @@ limitations under the License. #include "tensorflow/compiler/jit/build_xla_ops_pass.h" #include "tensorflow/compiler/jit/clone_constants_for_better_clustering.h" #include "tensorflow/compiler/jit/cluster_scoping_pass.h" -#include "tensorflow/compiler/jit/disable_functional_ops_lowering_for_xla_pass.h" #include "tensorflow/compiler/jit/encapsulate_subgraphs_pass.h" #include "tensorflow/compiler/jit/encapsulate_xla_computations_pass.h" #include "tensorflow/compiler/jit/force_xla_constants_on_host_pass.h" @@ -30,9 +29,6 @@ namespace tensorflow { // PRE_PLACEMENT passes: -REGISTER_OPTIMIZATION(OptimizationPassRegistry::PRE_PLACEMENT, 5, - DisableFunctionalOpsLoweringForXlaPass); - // EncapsulateXlaComputationsPass rewrites computations generated by the // xla.compile() Python code into XlaLaunch nodes. REGISTER_OPTIMIZATION(OptimizationPassRegistry::PRE_PLACEMENT, 36, diff --git a/tensorflow/compiler/tests/jit_test.py b/tensorflow/compiler/tests/jit_test.py index 80f8bdd57c51a9..32787da0b9ecc7 100644 --- a/tensorflow/compiler/tests/jit_test.py +++ b/tensorflow/compiler/tests/jit_test.py @@ -642,34 +642,6 @@ def CompiledFunction(x): self.assertTrue(InLabels(RunMetadataLabels(run_metadata), "_XlaRun")) -class AutoClusteringWhileLoopTest(test.TestCase): - """Regression test for auto-clustering silently dropping while_loop output. - - Auto-clustering used to lower functional control flow ops (While/If/Case) - to Switch/Merge before MarkForCompilationPass ran, so only fragments of a - loop body ended up in a cluster instead of the whole While op. That could - silently produce wrong results for some loop bodies, unlike - jit_compile=True which never lowers these ops in the first place. - """ - - def testWhileLoopOutputIsCorrectUnderAutoClustering(self): - config = NoRewriteSessionConfig() - config.graph_options.optimizer_options.global_jit_level = ( - config_pb2.OptimizerOptions.ON_2 - ) - - with session_lib.Session(config=config) as sess: - x = array_ops.placeholder(dtypes.float32) - c = lambda i, _: math_ops.less(i, 5) - b = lambda i, x: (i + 1, x + 1.0) - _, result = while_loop.while_loop(c, b, (constant_op.constant(0), x)) - - output = test_utils.RunWithWarmup( - sess, result, {x: np.array([1.0, 2.0, 3.0], dtype=np.float32)} - ) - self.assertAllClose(output, np.array([6.0, 7.0, 8.0], dtype=np.float32)) - - if __name__ == "__main__": os.environ["TF_XLA_FLAGS"] = ("--tf_xla_enable_lazy_compilation=true " + os.environ.get("TF_XLA_FLAGS", "")) From 43f829e7056c8e8ef69a53c16e8a1e9b60184ae2 Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Thu, 27 Aug 2026 00:43:46 -0700 Subject: [PATCH 27/32] Update platforms_config.patch to match new LLVM config.bzl An upstream LLVM change added Emscripten support to config.bzl, shifting line numbers and adding `@platforms//os:emscripten` to the context of `backtrace_defines`. This caused patch application to fail during OSS builds. PiperOrigin-RevId: 971781235 --- third_party/xla/third_party/llvm/platforms_config.patch | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/third_party/xla/third_party/llvm/platforms_config.patch b/third_party/xla/third_party/llvm/platforms_config.patch index 40af483204777c..995624e06d820e 100644 --- a/third_party/xla/third_party/llvm/platforms_config.patch +++ b/third_party/xla/third_party/llvm/platforms_config.patch @@ -15,8 +15,9 @@ diff --git a/utils/bazel/llvm-project-overlay/llvm/config.bzl b/utils/bazel/llvm-project-overlay/llvm/config.bzl --- a/utils/bazel/llvm-project-overlay/llvm/config.bzl +++ b/utils/bazel/llvm-project-overlay/llvm/config.bzl -@@ -49,13 +49,12 @@ +@@ -69,14 +69,13 @@ backtrace_defines = select({ + "@platforms//os:emscripten": [], "@platforms//os:windows": [], - "@llvm//platforms/config:musl": [], "//conditions:default": [ From c85ab3fa08ed3292a44bf5e92475dc702f159649 Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Thu, 27 Aug 2026 03:13:11 -0700 Subject: [PATCH 28/32] Add pointer_union patch for tf_runtime to fix compile error Upstream LLVM removed PointerUnion::get() and is(). This patches tf_runtime to use LLVM cast() in basic_kernels.cc for OSS builds. PiperOrigin-RevId: 971845980 --- third_party/tf_runtime/pointer_union.patch | 17 +++++++++++++++++ third_party/tf_runtime/workspace.bzl | 5 ++++- 2 files changed, 21 insertions(+), 1 deletion(-) create mode 100644 third_party/tf_runtime/pointer_union.patch diff --git a/third_party/tf_runtime/pointer_union.patch b/third_party/tf_runtime/pointer_union.patch new file mode 100644 index 00000000000000..334a96ce722eb0 --- /dev/null +++ b/third_party/tf_runtime/pointer_union.patch @@ -0,0 +1,17 @@ +diff --git a/lib/basic_kernels/opdefs/basic_kernels.cc b/lib/basic_kernels/opdefs/basic_kernels.cc +--- a/lib/basic_kernels/opdefs/basic_kernels.cc ++++ b/lib/basic_kernels/opdefs/basic_kernels.cc +@@ -137,11 +137,11 @@ + void CallOp::setCalleeFromCallable(CallInterfaceCallable callee) { + // Direct call. + if (FlatSymbolRefAttr calleeAttr = getCalleeAttr()) { +- auto symRef = callee.get(); ++ auto symRef = cast(callee); + return setCalleeAttr(cast(symRef)); + } + // Indirect call, callee Value is the first operand. +- return setOperand(0, callee.get()); ++ return setOperand(0, cast(callee)); + } + + //===----------------------------------------------------------------------===// diff --git a/third_party/tf_runtime/workspace.bzl b/third_party/tf_runtime/workspace.bzl index 7bc99867ef45b5..33d518ba0ac0db 100644 --- a/third_party/tf_runtime/workspace.bzl +++ b/third_party/tf_runtime/workspace.bzl @@ -31,5 +31,8 @@ def repo(): urls = tf_mirror_urls("https://github.com/tensorflow/runtime/archive/{commit}.tar.gz".format(commit = TFRT_COMMIT)), # A patch file can be provided for atomic commits to both TF and TFRT. # The job that bumps the TFRT_COMMIT also resets patch_file to 'None'. - patch_file = ["//third_party/tf_runtime:f16_attr.patch"], + patch_file = [ + "//third_party/tf_runtime:f16_attr.patch", + "//third_party/tf_runtime:pointer_union.patch", + ], ) From 1d7dd3d178aea65d0b7c4604a06653ca403f8aac Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Thu, 27 Aug 2026 03:46:53 -0700 Subject: [PATCH 29/32] Delete computation_placer_hdr target. All targets that previously depended on this target now either depend on computation_placer or device_assignment targets. PiperOrigin-RevId: 971858217 --- third_party/xla/xla/service/BUILD | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD index 25c6fdf0f4daca..152db83f4beeb1 100644 --- a/third_party/xla/xla/service/BUILD +++ b/third_party/xla/xla/service/BUILD @@ -3701,16 +3701,6 @@ xla_cc_test( ], ) -cc_library( - name = "computation_placer_hdr", - hdrs = ["computation_placer.h"], - deps = [ - ":device_assignment", - "//xla/stream_executor:platform_id", - "@com_google_absl//absl/status:statusor", - ], -) - cc_library( name = "human_readable_profile_builder", srcs = ["human_readable_profile_builder.cc"], From c02c66017633e4e56b0fc389d143f2cae577f84f Mon Sep 17 00:00:00 2001 From: Alexander Lyashuk Date: Thu, 27 Aug 2026 03:53:11 -0700 Subject: [PATCH 30/32] Reverts 73089e3228ff1893585f947c243ec426db9a2b0f PiperOrigin-RevId: 971860739 --- .../xla/xla/backends/gpu/transforms/BUILD | 39 -- .../gpu/transforms/alias_in_place_outputs.cc | 208 --------- .../gpu/transforms/alias_in_place_outputs.h | 65 --- .../transforms/alias_in_place_outputs_test.cc | 434 ------------------ third_party/xla/xla/service/gpu/BUILD | 1 - .../xla/xla/service/gpu/gpu_compiler.cc | 5 +- .../pre_scheduling_copy_insertion_pipeline.cc | 6 +- .../pre_scheduling_copy_insertion_pipeline.h | 7 +- 8 files changed, 4 insertions(+), 761 deletions(-) delete mode 100644 third_party/xla/xla/backends/gpu/transforms/alias_in_place_outputs.cc delete mode 100644 third_party/xla/xla/backends/gpu/transforms/alias_in_place_outputs.h delete mode 100644 third_party/xla/xla/backends/gpu/transforms/alias_in_place_outputs_test.cc diff --git a/third_party/xla/xla/backends/gpu/transforms/BUILD b/third_party/xla/xla/backends/gpu/transforms/BUILD index f7ee78cd9db44b..864d7e0fe5126b 100644 --- a/third_party/xla/xla/backends/gpu/transforms/BUILD +++ b/third_party/xla/xla/backends/gpu/transforms/BUILD @@ -182,45 +182,6 @@ xla_cc_test( ], ) -cc_library( - name = "alias_in_place_outputs", - srcs = ["alias_in_place_outputs.cc"], - hdrs = ["alias_in_place_outputs.h"], - deps = [ - "//xla:shape_util", - "//xla/hlo/analysis:hlo_reachability", - "//xla/hlo/analysis:indexing_analysis", - "//xla/hlo/analysis:symbolic_map", - "//xla/hlo/ir:hlo", - "//xla/hlo/pass:hlo_pass", - "//xla/hlo/utils:hlo_traversal", - "//xla/service/gpu:backend_configs_cc", - "//xla/service/gpu:cublas_cudnn", - "//xla/service/gpu:ir_emission_utils", - "@com_google_absl//absl/algorithm:container", - "@com_google_absl//absl/container:flat_hash_set", - "@com_google_absl//absl/log", - "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/strings:string_view", - "@llvm-project//mlir:IR", - ], -) - -xla_cc_test( - name = "alias_in_place_outputs_test", - srcs = ["alias_in_place_outputs_test.cc"], - deps = [ - ":alias_in_place_outputs", - "//xla:shape_util", - "//xla/hlo/ir:hlo", - "//xla/hlo/testlib:hlo_hardware_independent_test_base", - "@com_google_absl//absl/log", - "@com_google_absl//absl/strings:string_view", - "@com_google_googletest//:gtest_main", - "@llvm-project//mlir:IR", - ], -) - cc_library( name = "async_wrapper", srcs = ["async_wrapper.cc"], diff --git a/third_party/xla/xla/backends/gpu/transforms/alias_in_place_outputs.cc b/third_party/xla/xla/backends/gpu/transforms/alias_in_place_outputs.cc deleted file mode 100644 index 2d87ff77af0b93..00000000000000 --- a/third_party/xla/xla/backends/gpu/transforms/alias_in_place_outputs.cc +++ /dev/null @@ -1,208 +0,0 @@ -/* Copyright 2026 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/backends/gpu/transforms/alias_in_place_outputs.h" - -#include -#include -#include - -#include "absl/algorithm/container.h" -#include "absl/container/flat_hash_set.h" -#include "absl/log/log.h" -#include "absl/status/statusor.h" -#include "absl/strings/string_view.h" -#include "mlir/IR/MLIRContext.h" -#include "xla/hlo/analysis/hlo_reachability.h" -#include "xla/hlo/analysis/indexing_analysis.h" -#include "xla/hlo/analysis/indexing_map.h" -#include "xla/hlo/analysis/symbolic_expr.h" -#include "xla/hlo/ir/hlo_casting_utils.h" -#include "xla/hlo/ir/hlo_computation.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/hlo/ir/hlo_instructions.h" -#include "xla/hlo/ir/hlo_opcode.h" -#include "xla/hlo/utils/hlo_traversal.h" -#include "xla/service/gpu/backend_configs.pb.h" -#include "xla/service/gpu/cublas_cudnn.h" -#include "xla/service/gpu/ir_emission_utils.h" -#include "xla/shape.h" -#include "xla/shape_util.h" - -namespace xla::gpu { -namespace { - -bool IsTritonFusion(const HloFusionInstruction& fusion) { - auto config = fusion.backend_config(); - if (!config.ok()) { - return false; - } - absl::string_view kind = config->fusion_backend_config().kind(); - return kind == kTritonFusionKind || kind == kTritonGemmFusionKind || - kind == kTritonNestedGemmFusionKind; -} - -bool IsSupportedCublasLtMatmul(const HloInstruction& hlo) { - return IsCublasLtMatmul(hlo) || IsCublasLtMatmulF8(hlo) || - IsCublasLtMatmulMx(hlo) || IsCublasLtGroupedMatmul(hlo); -} - -// Returns true if the composed output→input indexing map for the given operand -// of the fusion is an identity, i.e. output element [i,j,...] reads input -// element [i,j,...]. This covers pure elementwise chains and bitcasts that -// cancel. Operands read through multiple data-flow paths are rejected. -bool HasIdentityIndexing(const GroupedByOpIndexing& grouped, - const HloInstruction* operand) { - auto it = grouped.find(operand); - if (it == grouped.end() || it->second.size() != 1) { - return false; - } - const IndexingMap& map = it->second.begin()->map(); - return map.GetSymbolicMap().IsIdentity() && map.GetRangeVars().empty() && - map.GetSymbolicConstraints().empty(); -} - -// Returns true if overwriting `operand` in place avoids rather than forces a -// copy: `operand` must be a writable intermediate (not a parameter, constant, -// or root) that `user` reads exactly once, and every other user of `operand` -// must precede `user` in the graph (e.g. a residual read before the op -// overwrites it). -bool BeneficialToAlias(const HloInstruction* operand, - const HloInstruction* user, - const HloReachabilityMap& reachability) { - if (operand->opcode() == HloOpcode::kParameter || - operand->opcode() == HloOpcode::kConstant || operand->IsRoot()) { - return false; - } - if (absl::c_count(user->operands(), operand) != 1) { - // Don't alias if the same buffer is used multiple times in the user. - return false; - } - return absl::c_all_of(operand->users(), [&](const HloInstruction* other) { - return other == user || reachability.IsReachable(other, user); - }); -} - -bool AliasFusion(HloFusionInstruction& fusion, - const HloReachabilityMap& reachability, - mlir::MLIRContext& mlir_ctx) { - if (!fusion.shape().IsArray() || - !fusion.output_to_operand_aliasing().empty()) { - return false; - } - std::unique_ptr fusion_adaptor = - HloFusionAdaptor::ForInstruction(&fusion); - GroupedByOpIndexing grouped = ComputeGroupedOutputToInputIndexing( - *fusion_adaptor, fusion_adaptor->GetRoots()[0], &mlir_ctx); - const HloInstruction* root = fusion.fused_expression_root(); - for (int64_t i = 0; i < fusion.operand_count(); ++i) { - const HloInstruction* operand = fusion.operand(i); - // Aliasing requires matching element type and dimensions (layouts may - // differ), as enforced by the HLO verifier. - if (!ShapeUtil::Compatible(operand->shape(), root->shape())) { - continue; - } - if (!HasIdentityIndexing(grouped, operand)) { - continue; - } - if (!BeneficialToAlias(operand, &fusion, reachability)) { - continue; - } - VLOG(2) << "Aliasing output of " << fusion.name() << " to operand " << i - << " (" << operand->name() << ")"; - fusion.set_output_to_operand_aliasing( - {{/*output_index=*/{}, {/*operand_number=*/i, /*operand_index=*/{}}}}); - return true; - } - return false; -} - -bool AliasCublasLtMatmul(HloCustomCallInstruction& custom_call, - const HloReachabilityMap& reachability) { - if (!custom_call.output_to_operand_aliasing().empty()) { - return false; - } - - const int64_t bias_idx = IsCublasLtGroupedMatmul(custom_call) ? 3 : 2; - if (custom_call.operand_count() <= bias_idx) { - return false; - } - - auto config = custom_call.backend_config(); - if (!config.ok() || config->gemm_backend_config().beta() == 0.0) { - return false; - } - - // cuBLASLt returns either the result array or a (result, scratch) tuple. - const Shape& shape = custom_call.shape(); - ShapeIndex output_index; - const Shape* output_shape = &shape; - if (shape.IsTuple()) { - if (shape.tuple_shapes().empty()) { - return false; - } - output_index = ShapeIndex{0}; - output_shape = &shape.tuple_shapes(0); - } else if (!shape.IsArray()) { - return false; - } - - const HloInstruction* bias = custom_call.operand(bias_idx); - if (!ShapeUtil::Equal(bias->shape(), *output_shape)) { - return false; - } - if (!BeneficialToAlias(bias, &custom_call, reachability)) { - return false; - } - - VLOG(2) << "Aliasing output " << output_index.ToString() << " of " - << custom_call.name() << " to operand " << bias_idx << " (" - << bias->name() << ")"; - custom_call.set_output_to_operand_aliasing( - {{output_index, {/*operand_number=*/bias_idx, /*operand_index=*/{}}}}); - return true; -} - -} // namespace - -absl::StatusOr AliasInPlaceOutputs::RunImpl( - HloModule* module, - const absl::flat_hash_set& execution_threads) { - bool changed = false; - RegisterSymbolicExprStorage(mlir_context_); - for (HloComputation* computation : - module->MakeNonfusionComputations(execution_threads)) { - std::unique_ptr reachability; - // Only build the reachability map if needed, reuse if already built. - const auto reachability_of = [&]() -> const HloReachabilityMap& { - if (reachability == nullptr) { - reachability = HloReachabilityMap::Build(computation); - } - return *reachability; - }; - for (HloInstruction* instr : computation->instructions()) { - if (auto* fusion = DynCast(instr); - fusion != nullptr && IsTritonFusion(*fusion)) { - changed |= AliasFusion(*fusion, reachability_of(), *mlir_context_); - } else if (auto* call = DynCast(instr); - call != nullptr && IsSupportedCublasLtMatmul(*call)) { - changed |= AliasCublasLtMatmul(*call, reachability_of()); - } - } - } - return changed; -} - -} // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/transforms/alias_in_place_outputs.h b/third_party/xla/xla/backends/gpu/transforms/alias_in_place_outputs.h deleted file mode 100644 index f833678195ae27..00000000000000 --- a/third_party/xla/xla/backends/gpu/transforms/alias_in_place_outputs.h +++ /dev/null @@ -1,65 +0,0 @@ -/* Copyright 2026 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef XLA_BACKENDS_GPU_TRANSFORMS_ALIAS_IN_PLACE_OUTPUTS_H_ -#define XLA_BACKENDS_GPU_TRANSFORMS_ALIAS_IN_PLACE_OUTPUTS_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/pass/hlo_pass_interface.h" - -namespace mlir { -class MLIRContext; -} // namespace mlir - -namespace xla::gpu { - -// Annotates GPU operations whose output buffer can overwrite an operand buffer -// in place with `output_to_operand_aliasing`, eliminating copies. -// -// Supported operations: -// * Triton fusions (GEMMs, elementwise): an operand is aliased when it -// reaches the fusion root through a single data-flow path consisting of -// pure elementwise ops or bitcasts that cancel, and its shape is -// compatible with the root. -// * cuBLASLt custom calls (`__cublas$lt$matmul`, `__cublas$lt$matmul$f8`, -// `__cublas$lt$matmul$mx`, `__cublas$lt$groupedMatmul`): the bias/C -// operand is aliased when `beta != 0` and its shape matches the output. -// -// In both cases, the aliased operand must be a writable intermediate whose -// only other users (if any) precede this operation in the dataflow graph. -// -// Must run before copy insertion. -class AliasInPlaceOutputs : public HloModulePass { - public: - explicit AliasInPlaceOutputs(mlir::MLIRContext* mlir_context) - : mlir_context_(mlir_context) {} - - absl::string_view name() const override { return "alias_in_place_outputs"; } - - protected: - absl::StatusOr RunImpl( - HloModule* module, - const absl::flat_hash_set& execution_threads) override; - - private: - mlir::MLIRContext* mlir_context_; -}; - -} // namespace xla::gpu - -#endif // XLA_BACKENDS_GPU_TRANSFORMS_ALIAS_IN_PLACE_OUTPUTS_H_ diff --git a/third_party/xla/xla/backends/gpu/transforms/alias_in_place_outputs_test.cc b/third_party/xla/xla/backends/gpu/transforms/alias_in_place_outputs_test.cc deleted file mode 100644 index 00c67074d1bc22..00000000000000 --- a/third_party/xla/xla/backends/gpu/transforms/alias_in_place_outputs_test.cc +++ /dev/null @@ -1,434 +0,0 @@ -/* Copyright 2026 The OpenXLA Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "xla/backends/gpu/transforms/alias_in_place_outputs.h" - -#include -#include -#include -#include - -#include -#include -#include "absl/log/log.h" -#include "absl/strings/string_view.h" -#include "mlir/IR/MLIRContext.h" -#include "xla/hlo/ir/hlo_casting_utils.h" -#include "xla/hlo/ir/hlo_instruction.h" -#include "xla/hlo/ir/hlo_instructions.h" -#include "xla/hlo/ir/hlo_module.h" -#include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" -#include "xla/shape_util.h" - -namespace xla::gpu { -namespace { - -using ::testing::ElementsAre; -using ::testing::Pair; - -class AliasInPlaceOutputsTest : public HloHardwareIndependentTestBase { - protected: - using AliasingList = - std::vector>>; - - // Output-to-operand aliasing of the single candidate instruction - // (Triton fusion or cuBLASLt custom call) in the entry computation. - const AliasingList& GetAliasing(const HloModule& module) { - for (const HloInstruction* instr : - module.entry_computation()->instructions()) { - if (const auto* fusion = DynCast(instr)) { - return fusion->output_to_operand_aliasing(); - } - if (const auto* call = DynCast(instr)) { - return call->output_to_operand_aliasing(); - } - } - LOG(FATAL) << "no candidate instruction in entry computation"; - } - - void ExpectAlias(absl::string_view hlo, int64_t operand, - ShapeIndex output = {}) { - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo)); - ASSERT_OK_AND_ASSIGN(bool changed, - AliasInPlaceOutputs(&mlir_context_).Run(module.get())); - EXPECT_TRUE(changed); - EXPECT_THAT(GetAliasing(*module), - ElementsAre(Pair(output, Pair(operand, ShapeIndex{})))); - } - - void ExpectNoAlias(absl::string_view hlo) { - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo)); - ASSERT_OK_AND_ASSIGN(bool changed, - AliasInPlaceOutputs(&mlir_context_).Run(module.get())); - EXPECT_FALSE(changed); - EXPECT_TRUE(GetAliasing(*module).empty()); - } - - mlir::MLIRContext mlir_context_; -}; - -// --- Triton Fusion Tests --- - -// dot(a, b) + x with x a dead intermediate => x is aliased to the output. -TEST_F(AliasInPlaceOutputsTest, TritonAliasesDeadBiasOperand) { - ExpectAlias(R"( - HloModule m - - triton { - a = f32[128,64] parameter(0) - b = f32[64,128] parameter(1) - x = f32[128,128] parameter(2) - dot = f32[128,128] dot(a, b), - lhs_contracting_dims={1}, rhs_contracting_dims={0} - ROOT add = f32[128,128] add(dot, x) - } - - ENTRY e { - a = f32[128,64] parameter(0) - b = f32[64,128] parameter(1) - c = f32[128,128] parameter(2) - x = f32[128,128] add(c, c) - ROOT fusion = f32[128,128] fusion(a, b, x), kind=kCustom, calls=triton, - backend_config={"fusion_backend_config":{"kind":"__triton_gemm"}} - })", - /*operand=*/2); -} - -// Post-autotuning Triton GEMMs are converted to __triton_nested_gemm_fusion. -TEST_F(AliasInPlaceOutputsTest, TritonAliasesPostAutotuningNestedGemmFusion) { - ExpectAlias(R"( - HloModule m - - triton { - a = f32[128,64] parameter(0) - b = f32[64,128] parameter(1) - x = f32[128,128] parameter(2) - dot = f32[128,128] dot(a, b), - lhs_contracting_dims={1}, rhs_contracting_dims={0} - ROOT add = f32[128,128] add(dot, x) - } - - ENTRY e { - a = f32[128,64] parameter(0) - b = f32[64,128] parameter(1) - c = f32[128,128] parameter(2) - x = f32[128,128] add(c, c) - ROOT fusion = f32[128,128] fusion(a, b, x), kind=kCustom, calls=triton, - backend_config={"fusion_backend_config":{"kind":"__triton_nested_gemm_fusion"}} - })", - /*operand=*/2); -} - -// The bias reaches the root through a non-identity access (transpose), so it is -// not safe to alias even though shapes match. -TEST_F(AliasInPlaceOutputsTest, TritonDoesNotAliasNonIdentityAccess) { - ExpectNoAlias(R"( - HloModule m - - triton { - a = f32[128,64] parameter(0) - b = f32[64,128] parameter(1) - x = f32[128,128] parameter(2) - xt = f32[128,128] transpose(x), dimensions={1,0} - dot = f32[128,128] dot(a, b), - lhs_contracting_dims={1}, rhs_contracting_dims={0} - ROOT add = f32[128,128] add(dot, xt) - } - - ENTRY e { - a = f32[128,64] parameter(0) - b = f32[64,128] parameter(1) - c = f32[128,128] parameter(2) - x = f32[128,128] add(c, c) - ROOT fusion = f32[128,128] fusion(a, b, x), kind=kCustom, calls=triton, - backend_config={"fusion_backend_config":{"kind":"__triton_gemm"}} - })"); -} - -// The bias is a plain entry parameter (read-only input) => not beneficial. -TEST_F(AliasInPlaceOutputsTest, TritonDoesNotAliasParameterOperand) { - ExpectNoAlias(R"( - HloModule m - - triton { - a = f32[128,64] parameter(0) - b = f32[64,128] parameter(1) - x = f32[128,128] parameter(2) - dot = f32[128,128] dot(a, b), - lhs_contracting_dims={1}, rhs_contracting_dims={0} - ROOT add = f32[128,128] add(dot, x) - } - - ENTRY e { - a = f32[128,64] parameter(0) - b = f32[64,128] parameter(1) - x = f32[128,128] parameter(2) - ROOT fusion = f32[128,128] fusion(a, b, x), kind=kCustom, calls=triton, - backend_config={"fusion_backend_config":{"kind":"__triton_gemm"}} - })"); -} - -// The bias has another user that is independent => not safe to alias. -TEST_F(AliasInPlaceOutputsTest, TritonDoesNotAliasIndependentLiveOperand) { - ExpectNoAlias(R"( - HloModule m - - triton { - a = f32[128,64] parameter(0) - b = f32[64,128] parameter(1) - x = f32[128,128] parameter(2) - dot = f32[128,128] dot(a, b), - lhs_contracting_dims={1}, rhs_contracting_dims={0} - ROOT add = f32[128,128] add(dot, x) - } - - ENTRY e { - a = f32[128,64] parameter(0) - b = f32[64,128] parameter(1) - c = f32[128,128] parameter(2) - x = f32[128,128] add(c, c) - fusion = f32[128,128] fusion(a, b, x), kind=kCustom, calls=triton, - backend_config={"fusion_backend_config":{"kind":"__triton_gemm"}} - ROOT t = (f32[128,128], f32[128,128]) tuple(fusion, x) - })"); -} - -// The bias has another user, but that user is a predecessor of the fusion -// (residual pattern: a = norm(x); fusion = dot(a, b) + x). Safe to alias! -TEST_F(AliasInPlaceOutputsTest, - TritonAliasesResidualOperandWithPredecessorUser) { - ExpectAlias(R"( - HloModule m - - triton { - a = f32[128,64] parameter(0) - b = f32[64,128] parameter(1) - x = f32[128,128] parameter(2) - dot = f32[128,128] dot(a, b), - lhs_contracting_dims={1}, rhs_contracting_dims={0} - ROOT add = f32[128,128] add(dot, x) - } - - ENTRY e { - c = f32[128,128] parameter(0) - b = f32[64,128] parameter(1) - x = f32[128,128] add(c, c) - slice = f32[128,64] slice(x), slice={[0:128], [0:64]} - ROOT fusion = f32[128,128] fusion(slice, b, x), kind=kCustom, calls=triton, - backend_config={"fusion_backend_config":{"kind":"__triton_gemm"}} - })", - /*operand=*/2); -} - -// --- cuBLASLt Custom Call Tests --- - -// __cublas$lt$matmul with beta=1 and dead bias operand => aliases output {0} to -// operand {2, {}}. -TEST_F(AliasInPlaceOutputsTest, CublasLtAliasesDeadBiasOperand) { - ExpectAlias(R"( - HloModule m - - ENTRY e { - p0 = f32[128,64] parameter(0) - p1 = f32[64,128] parameter(1) - p2 = f32[128,128] parameter(2) - bias = f32[128,128] add(p2, p2) - ROOT matmul = (f32[128,128], s8[4096]) custom-call(p0, p1, bias), - custom_call_target="__cublas$lt$matmul", - backend_config={"gemm_backend_config":{"beta":1.0}} - })", - /*operand=*/2, /*output=*/ShapeIndex{0}); -} - -// __cublas$lt$matmul with beta=0 => no aliasing. -TEST_F(AliasInPlaceOutputsTest, CublasLtDoesNotAliasWhenBetaIsZero) { - ExpectNoAlias(R"( - HloModule m - - ENTRY e { - p0 = f32[128,64] parameter(0) - p1 = f32[64,128] parameter(1) - p2 = f32[128,128] parameter(2) - bias = f32[128,128] add(p2, p2) - ROOT matmul = (f32[128,128], s8[4096]) custom-call(p0, p1, bias), - custom_call_target="__cublas$lt$matmul", - backend_config={"gemm_backend_config":{"beta":0.0}} - })"); -} - -// __cublas$lt$matmul where bias is a parameter => no aliasing. -TEST_F(AliasInPlaceOutputsTest, CublasLtDoesNotAliasParameterOperand) { - ExpectNoAlias(R"( - HloModule m - - ENTRY e { - p0 = f32[128,64] parameter(0) - p1 = f32[64,128] parameter(1) - p2 = f32[128,128] parameter(2) - ROOT matmul = (f32[128,128], s8[4096]) custom-call(p0, p1, p2), - custom_call_target="__cublas$lt$matmul", - backend_config={"gemm_backend_config":{"beta":1.0}} - })"); -} - -// cuBLASLt in a residual chain (AlphaFold3 pattern): -// bias is used by an intermediate layer (e.g. norm/slice) that produces LHS of -// matmul. Because the other user precedes matmul, aliasing is safe! -TEST_F(AliasInPlaceOutputsTest, CublasLtAliasesResidualChain) { - ExpectAlias(R"( - HloModule m - - ENTRY e { - p0 = f32[128,128] parameter(0) - p1 = f32[64,128] parameter(1) - bias = f32[128,128] add(p0, p0) - lhs = f32[128,64] slice(bias), slice={[0:128], [0:64]} - matmul = (f32[128,128], s8[4096]) custom-call(lhs, p1, bias), - custom_call_target="__cublas$lt$matmul", - backend_config={"gemm_backend_config":{"beta":1.0}} - ROOT gte = f32[128,128] get-tuple-element(matmul), index=0 - })", - /*operand=*/2, /*output=*/ShapeIndex{0}); -} - -// cuBLASLt where other user does NOT precede matmul => no aliasing. -TEST_F(AliasInPlaceOutputsTest, CublasLtDoesNotAliasIndependentUser) { - ExpectNoAlias(R"( - HloModule m - - ENTRY e { - p0 = f32[128,64] parameter(0) - p1 = f32[64,128] parameter(1) - p2 = f32[128,128] parameter(2) - bias = f32[128,128] add(p2, p2) - matmul = (f32[128,128], s8[4096]) custom-call(p0, p1, bias), - custom_call_target="__cublas$lt$matmul", - backend_config={"gemm_backend_config":{"beta":1.0}} - ROOT t = ((f32[128,128], s8[4096]), f32[128,128]) tuple(matmul, bias) - })"); -} - -// cuBLASLt FP8 (__cublas$lt$matmul$f8) with beta=1. -TEST_F(AliasInPlaceOutputsTest, CublasLtF8AliasesDeadBiasOperand) { - ExpectAlias(R"( - HloModule m - - ENTRY e { - p0 = f8e4m3fn[128,64] parameter(0) - p1 = f8e4m3fn[64,128] parameter(1) - p2 = f32[128,128] parameter(2) - bias = f32[128,128] add(p2, p2) - ROOT matmul = (f32[128,128], s8[4096]) custom-call(p0, p1, bias), - custom_call_target="__cublas$lt$matmul$f8", - backend_config={"gemm_backend_config":{"beta":1.0}} - })", - /*operand=*/2, /*output=*/ShapeIndex{0}); -} - -// Grouped GEMM (__cublas$lt$groupedMatmul) with beta=1 (bias at operand index -// 3). -TEST_F(AliasInPlaceOutputsTest, CublasLtGroupedMatmulAliasesBiasOperand) { - ExpectAlias(R"( - HloModule m - - ENTRY e { - p0 = f32[128,64] parameter(0) - p1 = f32[64,128] parameter(1) - sizes = s32[2] parameter(2) - p3 = f32[128,128] parameter(3) - bias = f32[128,128] add(p3, p3) - ROOT matmul = (f32[128,128], s8[4096]) custom-call(p0, p1, sizes, bias), - custom_call_target="__cublas$lt$groupedMatmul", - backend_config={"gemm_backend_config":{"beta":1.0}} - })", - /*operand=*/3, /*output=*/ShapeIndex{0}); -} - -// Block-scaled MX matmul (__cublas$lt$matmul$mx) with beta=1. -TEST_F(AliasInPlaceOutputsTest, CublasLtMatmulMxAliasesDeadBiasOperand) { - ExpectAlias(R"( - HloModule m - - ENTRY e { - p0 = f8e4m3fn[128,64] parameter(0) - p1 = f8e4m3fn[64,128] parameter(1) - p2 = f32[128,128] parameter(2) - bias = f32[128,128] add(p2, p2) - ROOT matmul = (f32[128,128], s8[4096]) custom-call(p0, p1, bias), - custom_call_target="__cublas$lt$matmul$mx", - backend_config={"gemm_backend_config":{"beta":1.0}} - })", - /*operand=*/2, /*output=*/ShapeIndex{0}); -} - -// Triton fusion with tuple output is not supported for aliasing. -TEST_F(AliasInPlaceOutputsTest, TritonDoesNotAliasTupleFusion) { - ExpectNoAlias(R"( - HloModule m - - triton { - a = f32[128,64] parameter(0) - b = f32[64,128] parameter(1) - x = f32[128,128] parameter(2) - dot = f32[128,128] dot(a, b), - lhs_contracting_dims={1}, rhs_contracting_dims={0} - add = f32[128,128] add(dot, x) - ROOT t = (f32[128,128], f32[128,128]) tuple(add, dot) - } - - ENTRY e { - a = f32[128,64] parameter(0) - b = f32[64,128] parameter(1) - c = f32[128,128] parameter(2) - x = f32[128,128] add(c, c) - ROOT fusion = (f32[128,128], f32[128,128]) fusion(a, b, x), kind=kCustom, calls=triton, - backend_config={"fusion_backend_config":{"kind":"__triton_gemm"}} - })"); -} - -// cuBLASLt with constant bias operand => no aliasing. -TEST_F(AliasInPlaceOutputsTest, CublasLtDoesNotAliasConstantOperand) { - ExpectNoAlias(R"( - HloModule m - - ENTRY e { - p0 = f32[128,64] parameter(0) - p1 = f32[64,128] parameter(1) - c = f32[128,128] constant({...}) - ROOT matmul = (f32[128,128], s8[4096]) custom-call(p0, p1, c), - custom_call_target="__cublas$lt$matmul", - backend_config={"gemm_backend_config":{"beta":1.0}} - })"); -} - -// cuBLASLt where bias shape does not match output shape (e.g. vector broadcast -// bias) => no aliasing. -TEST_F(AliasInPlaceOutputsTest, CublasLtDoesNotAliasMismatchedShapeBias) { - ExpectNoAlias(R"( - HloModule m - - ENTRY e { - p0 = f32[128,64] parameter(0) - p1 = f32[64,128] parameter(1) - p2 = f32[128] parameter(2) - bias = f32[128] add(p2, p2) - ROOT matmul = (f32[128,128], s8[4096]) custom-call(p0, p1, bias), - custom_call_target="__cublas$lt$matmul", - backend_config={"gemm_backend_config":{"beta":1.0}} - })"); -} - -} // namespace -} // namespace xla::gpu diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD index 156ae91d48e18d..1357ef9030f976 100644 --- a/third_party/xla/xla/service/gpu/BUILD +++ b/third_party/xla/xla/service/gpu/BUILD @@ -1743,7 +1743,6 @@ cc_library( deps = [ ":alias_info", "//xla:xla_proto_cc", - "//xla/backends/gpu/transforms:alias_in_place_outputs", "//xla/backends/gpu/transforms:alias_passthrough_params", "//xla/backends/gpu/transforms:copy_fusion", "//xla/backends/gpu/transforms:sanitize_constant_names", diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc index 5658324242ab93..0939d1c897195e 100644 --- a/third_party/xla/xla/service/gpu/gpu_compiler.cc +++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc @@ -1992,11 +1992,8 @@ absl::Status GpuCompiler::OptimizeHloModule( absl::Status GpuCompiler::RunPreSchedulingCopyInsertion( HloModule& hlo_module, const se::DeviceDescription& device_description, const GpuAliasInfo* alias_info) { - ABSL_ASSIGN_OR_RETURN(BorrowedMlirContext borrowed_context, - mlir_context_pool_.GetOrCreate()); - mlir::MLIRContext* mlir_context = borrowed_context->get(); return PreSchedulingCopyInsertionPipeline(hlo_module.config(), alias_info, - device_description, mlir_context) + device_description) .Run(&hlo_module, {HloInstruction::kMainExecutionThread}) .status(); } diff --git a/third_party/xla/xla/service/gpu/pre_scheduling_copy_insertion_pipeline.cc b/third_party/xla/xla/service/gpu/pre_scheduling_copy_insertion_pipeline.cc index 824e37e3e2da53..b8ac06a1546d5b 100644 --- a/third_party/xla/xla/service/gpu/pre_scheduling_copy_insertion_pipeline.cc +++ b/third_party/xla/xla/service/gpu/pre_scheduling_copy_insertion_pipeline.cc @@ -19,10 +19,10 @@ limitations under the License. #include #include -#include "xla/backends/gpu/transforms/alias_in_place_outputs.h" #include "xla/backends/gpu/transforms/alias_passthrough_params.h" #include "xla/backends/gpu/transforms/copy_fusion.h" #include "xla/backends/gpu/transforms/sanitize_constant_names.h" +#include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/pass/hlo_pass_pipeline.h" #include "xla/hlo/transforms/simplifiers/hlo_dce.h" #include "xla/service/copy_insertion.h" @@ -40,8 +40,7 @@ namespace gpu { HloPassPipeline PreSchedulingCopyInsertionPipeline( const HloModuleConfig& config, const GpuAliasInfo* alias_info, - const se::DeviceDescription& device_description, - mlir::MLIRContext* mlir_context) { + const se::DeviceDescription& device_description) { const DebugOptions& debug_options = config.debug_options(); // In some cases, we have to place the result of an instruction in a temporary @@ -68,7 +67,6 @@ HloPassPipeline PreSchedulingCopyInsertionPipeline( if (config.alias_passthrough_params()) { pipeline.AddPass(); } - pipeline.AddPass(mlir_context); pipeline.AddPass(alias_info); if (debug_options.xla_gpu_copy_insertion_use_region_analysis()) { diff --git a/third_party/xla/xla/service/gpu/pre_scheduling_copy_insertion_pipeline.h b/third_party/xla/xla/service/gpu/pre_scheduling_copy_insertion_pipeline.h index 7fdf8422aa0639..7b87575e35c0d5 100644 --- a/third_party/xla/xla/service/gpu/pre_scheduling_copy_insertion_pipeline.h +++ b/third_party/xla/xla/service/gpu/pre_scheduling_copy_insertion_pipeline.h @@ -21,10 +21,6 @@ limitations under the License. #include "xla/service/gpu/alias_info.h" #include "xla/stream_executor/device_description.h" -namespace mlir { -class MLIRContext; -} // namespace mlir - namespace xla { namespace gpu { @@ -32,8 +28,7 @@ namespace gpu { // This pipeline must run before scheduling to ensure correctness. HloPassPipeline PreSchedulingCopyInsertionPipeline( const HloModuleConfig& config, const GpuAliasInfo* alias_info, - const se::DeviceDescription& device_description, - mlir::MLIRContext* mlir_context); + const se::DeviceDescription& device_description); } // namespace gpu } // namespace xla From 4df202f7ca986d182b62bf750ad6bd3449ea3549 Mon Sep 17 00:00:00 2001 From: Christian Sigg Date: Thu, 27 Aug 2026 04:04:15 -0700 Subject: [PATCH 31/32] Add cuda_tile and tensor_ir repos. Configure XLA workspace import macros to support `@cuda_tile` and `@tensor_ir` dependencies in open-source XLA builds. PiperOrigin-RevId: 971865012 --- .../xla/third_party/cuda_tile/BUILD.bazel | 0 .../xla/third_party/cuda_tile/cuda_tile.BUILD | 388 ++++++++++++++ .../xla/third_party/cuda_tile/workspace.bzl | 16 + .../xla/third_party/tensor_ir/BUILD.bazel | 0 .../xla/third_party/tensor_ir/tensor_ir.BUILD | 480 ++++++++++++++++++ .../xla/third_party/tensor_ir/workspace.bzl | 16 + third_party/xla/workspace2.bzl | 4 + 7 files changed, 904 insertions(+) create mode 100644 third_party/xla/third_party/cuda_tile/BUILD.bazel create mode 100644 third_party/xla/third_party/cuda_tile/cuda_tile.BUILD create mode 100644 third_party/xla/third_party/cuda_tile/workspace.bzl create mode 100644 third_party/xla/third_party/tensor_ir/BUILD.bazel create mode 100644 third_party/xla/third_party/tensor_ir/tensor_ir.BUILD create mode 100644 third_party/xla/third_party/tensor_ir/workspace.bzl diff --git a/third_party/xla/third_party/cuda_tile/BUILD.bazel b/third_party/xla/third_party/cuda_tile/BUILD.bazel new file mode 100644 index 00000000000000..e69de29bb2d1d6 diff --git a/third_party/xla/third_party/cuda_tile/cuda_tile.BUILD b/third_party/xla/third_party/cuda_tile/cuda_tile.BUILD new file mode 100644 index 00000000000000..72a4d67e2d0297 --- /dev/null +++ b/third_party/xla/third_party/cuda_tile/cuda_tile.BUILD @@ -0,0 +1,388 @@ +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 = "cuda-tile-tblgen", + srcs = glob([ + "tools/cuda-tile-tblgen/*.cpp", + "tools/cuda-tile-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( + name = "CudaTileTdFiles", + srcs = glob(["include/cuda_tile/Dialect/CudaTile/IR/*.td"]), + includes = ["include"], + deps = [ + "@llvm-project//mlir:BuiltinDialectTdFiles", + "@llvm-project//mlir:ControlFlowInterfacesTdFiles", + "@llvm-project//mlir:FunctionInterfacesTdFiles", + "@llvm-project//mlir:InferTypeOpInterfaceTdFiles", + "@llvm-project//mlir:OpBaseTdFiles", + "@llvm-project//mlir:SideEffectInterfacesTdFiles", + "@llvm-project//mlir:ViewLikeInterfaceTdFiles", + ], +) + +gentbl_cc_library( + name = "CudaTileDialectIncGen", + tbl_outs = [ + ( + [ + "-gen-dialect-decls", + "-dialect=cuda_tile", + ], + "include/cuda_tile/Dialect/CudaTile/IR/Dialect.h.inc", + ), + ( + [ + "-gen-dialect-defs", + "-dialect=cuda_tile", + ], + "include/cuda_tile/Dialect/CudaTile/IR/Dialect.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/cuda_tile/Dialect/CudaTile/IR/Dialect.td", + deps = [":CudaTileTdFiles"], +) + +gentbl_cc_library( + name = "CudaTileInterfacesIncGen", + tbl_outs = [ + ( + ["-gen-attr-interface-decls"], + "include/cuda_tile/Dialect/CudaTile/IR/AttrInterfaces.h.inc", + ), + ( + ["-gen-attr-interface-defs"], + "include/cuda_tile/Dialect/CudaTile/IR/AttrInterfaces.cpp.inc", + ), + ( + ["-gen-type-interface-decls"], + "include/cuda_tile/Dialect/CudaTile/IR/TypeInterfaces.h.inc", + ), + ( + ["-gen-type-interface-defs"], + "include/cuda_tile/Dialect/CudaTile/IR/TypeInterfaces.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/cuda_tile/Dialect/CudaTile/IR/Interfaces.td", + deps = [":CudaTileTdFiles"], +) + +gentbl_cc_library( + name = "CudaTileTypesIncGen", + tbl_outs = [ + ( + [ + "-gen-typedef-decls", + "-typedefs-dialect=cuda_tile", + ], + "include/cuda_tile/Dialect/CudaTile/IR/Types.h.inc", + ), + ( + [ + "-gen-typedef-defs", + "-typedefs-dialect=cuda_tile", + ], + "include/cuda_tile/Dialect/CudaTile/IR/Types.cpp.inc", + ), + ( + [ + "-gen-type-constraint-decls", + "-dialect=cuda_tile", + ], + "include/cuda_tile/Dialect/CudaTile/IR/TypeConstraints.h.inc", + ), + ( + [ + "-gen-type-constraint-defs", + "-dialect=cuda_tile", + ], + "include/cuda_tile/Dialect/CudaTile/IR/TypeConstraints.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/cuda_tile/Dialect/CudaTile/IR/Types.td", + deps = [":CudaTileTdFiles"], +) + +gentbl_cc_library( + name = "CudaTileAttrDefsIncGen", + tbl_outs = [ + ( + [ + "-gen-attrdef-decls", + "-attrdefs-dialect=cuda_tile", + ], + "include/cuda_tile/Dialect/CudaTile/IR/AttrDefs.h.inc", + ), + ( + [ + "-gen-attrdef-defs", + "-attrdefs-dialect=cuda_tile", + ], + "include/cuda_tile/Dialect/CudaTile/IR/AttrDefs.cpp.inc", + ), + ( + [ + "-gen-enum-decls", + "-dialect=cuda_tile", + ], + "include/cuda_tile/Dialect/CudaTile/IR/Enums.h.inc", + ), + ( + [ + "-gen-enum-defs", + "-dialect=cuda_tile", + ], + "include/cuda_tile/Dialect/CudaTile/IR/Enums.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/cuda_tile/Dialect/CudaTile/IR/AttrDefs.td", + deps = [":CudaTileTdFiles"], +) + +gentbl_cc_library( + name = "CudaTileRemarksIncGen", + tbl_outs = [ + ( + ["-gen-enum-decls"], + "include/cuda_tile/Dialect/CudaTile/IR/TileIRRemarks.h.inc", + ), + ( + ["-gen-enum-defs"], + "include/cuda_tile/Dialect/CudaTile/IR/TileIRRemarks.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/cuda_tile/Dialect/CudaTile/IR/Remarks.td", + deps = [":CudaTileTdFiles"], +) + +gentbl_cc_library( + name = "CudaTileOpsIncGen", + tbl_outs = [ + ( + ["-gen-op-decls"], + "include/cuda_tile/Dialect/CudaTile/IR/Ops.h.inc", + ), + ( + ["-gen-op-defs"], + "include/cuda_tile/Dialect/CudaTile/IR/Ops.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/cuda_tile/Dialect/CudaTile/IR/Ops.td", + deps = [":CudaTileTdFiles"], +) + +gentbl_cc_library( + name = "CudaTileOpsCanonicalizationIncGen", + tbl_outs = [ + ( + ["-gen-rewriters"], + "lib/Dialect/CudaTile/IR/OpsCanonicalization.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "lib/Dialect/CudaTile/IR/OpsCanonicalization.td", + deps = [ + ":CudaTileTdFiles", + "@llvm-project//mlir:OpBaseTdFiles", + ], +) + +cc_library( + name = "CudaTileDialect", + srcs = glob( + ["lib/Dialect/CudaTile/IR/*.cpp"], + exclude = ["lib/Dialect/CudaTile/IR/CudaTileTesting.cpp"], + ), + hdrs = glob(["include/cuda_tile/Dialect/CudaTile/IR/*.h"]), + includes = [ + "include", + "lib/Dialect/CudaTile/IR", + ], + visibility = ["//visibility:public"], + deps = [ + ":CudaTileAttrDefsIncGen", + ":CudaTileDialectIncGen", + ":CudaTileInterfacesIncGen", + ":CudaTileOpsCanonicalizationIncGen", + ":CudaTileOpsIncGen", + ":CudaTileTypesIncGen", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:ArithDialect", + "@llvm-project//mlir:BytecodeOpInterface", + "@llvm-project//mlir:ControlFlowInterfaces", + "@llvm-project//mlir:FunctionInterfaces", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:InferTypeOpInterface", + "@llvm-project//mlir:InliningUtils", + "@llvm-project//mlir:SideEffectInterfaces", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:ViewLikeInterface", + ], +) + +td_library( + name = "CudaTileTransformsTdFiles", + srcs = ["include/cuda_tile/Dialect/CudaTile/Transforms/Passes.td"], + includes = ["include"], + deps = [ + "@llvm-project//mlir:PassBaseTdFiles", + ], +) + +gentbl_cc_library( + name = "CudaTileTransformsIncGen", + tbl_outs = [ + ( + [ + "-gen-pass-decls", + "-name=CudaTile", + ], + "include/cuda_tile/Dialect/CudaTile/Transforms/Passes.h.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/cuda_tile/Dialect/CudaTile/Transforms/Passes.td", + deps = [":CudaTileTransformsTdFiles"], +) + +cc_library( + name = "CudaTileTransforms", + srcs = glob(["lib/Dialect/CudaTile/Transforms/*.cpp"]), + hdrs = glob(["include/cuda_tile/Dialect/CudaTile/Transforms/*.h"]), + includes = ["include"], + visibility = ["//visibility:public"], + deps = [ + ":CudaTileDialect", + ":CudaTileTransformsIncGen", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:TransformUtils", + ], +) + +gentbl_cc_library( + name = "CudaTileBytecodeOpsIncGen", + tbl_outs = [ + ( + ["-gen-cuda-tile-bytecode"], + "include/cuda_tile/Bytecode/Writer/Bytecode.inc", + ), + ( + ["-gen-cuda-tile-bytecode-reader"], + "include/cuda_tile/Bytecode/Reader/BytecodeReader.inc", + ), + ], + tblgen = ":cuda-tile-tblgen", + td_file = "include/cuda_tile/Dialect/CudaTile/IR/Ops.td", + deps = [":CudaTileTdFiles"], +) + +gentbl_cc_library( + name = "CudaTileBytecodeOpcodesIncGen", + tbl_outs = [ + ( + ["-gen-cuda-tile-opcodes"], + "include/cuda_tile/Bytecode/Common/StaticOpcodes.inc", + ), + ], + tblgen = ":cuda-tile-tblgen", + td_file = "include/cuda_tile/Dialect/CudaTile/IR/BytecodeOpcodes.td", + deps = [":CudaTileTdFiles"], +) + +gentbl_cc_library( + name = "CudaTileBytecodeTypeIncGen", + tbl_outs = [ + ( + ["-gen-cuda-tile-type-bytecode"], + "include/cuda_tile/Bytecode/Writer/TypeBytecode.inc", + ), + ( + ["-gen-cuda-tile-type-bytecode-reader"], + "include/cuda_tile/Bytecode/Reader/TypeBytecodeReader.inc", + ), + ], + tblgen = ":cuda-tile-tblgen", + td_file = "include/cuda_tile/Dialect/CudaTile/IR/BytecodeTypeOpcodes.td", + deps = [":CudaTileTdFiles"], +) + +gentbl_cc_library( + name = "CudaTileBytecodeAttrIncGen", + tbl_outs = [ + ( + ["-gen-cuda-tile-attr-bytecode"], + "include/cuda_tile/Bytecode/Writer/AttrBytecode.inc", + ), + ], + tblgen = ":cuda-tile-tblgen", + td_file = "include/cuda_tile/Dialect/CudaTile/IR/BytecodeAttrOpcodes.td", + deps = [":CudaTileTdFiles"], +) + +cc_library( + name = "CudaTileBytecode", + srcs = glob([ + "lib/Bytecode/**/*.cpp", + "lib/Bytecode/**/*.h", + ]), + hdrs = glob([ + "include/cuda_tile/Bytecode/**/*.h", + ]), + includes = [ + "include", + "include/cuda_tile/Bytecode/Common", + "include/cuda_tile/Bytecode/Reader", + "include/cuda_tile/Bytecode/Writer", + ], + visibility = ["//visibility:public"], + deps = [ + ":CudaTileBytecodeAttrIncGen", + ":CudaTileBytecodeOpcodesIncGen", + ":CudaTileBytecodeOpsIncGen", + ":CudaTileBytecodeTypeIncGen", + ":CudaTileDialect", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:TranslateLib", + ], +) + +cc_library( + name = "CudaTileOptimizer", + srcs = ["lib/Dialect/CudaTile/Optimizer/CudaTileOptimizer.cpp"], + hdrs = ["include/cuda_tile/Dialect/CudaTile/Optimizer/CudaTileOptimizer.h"], + includes = ["include"], + visibility = ["//visibility:public"], + deps = [ + ":CudaTileBytecode", + ":CudaTileDialect", + ":CudaTileTransforms", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:Parser", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:Transforms", + ], +) diff --git a/third_party/xla/third_party/cuda_tile/workspace.bzl b/third_party/xla/third_party/cuda_tile/workspace.bzl new file mode 100644 index 00000000000000..7dcb4e0c479335 --- /dev/null +++ b/third_party/xla/third_party/cuda_tile/workspace.bzl @@ -0,0 +1,16 @@ +"""Provides the repository macro to import CUDA Tile IR.""" + +load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls") + +def repo(): + """Imports CUDA Tile IR.""" + CUDA_TILE_COMMIT = "af2417041cc939b87ef56d92cfdcf61737c5457e" + CUDA_TILE_SHA256 = "81597e49469171bf8fa7319fbd44ebe133001521f484589e3dd3fb3fad282dc0" + + tf_http_archive( + name = "cuda_tile", + build_file = "//third_party/cuda_tile:cuda_tile.BUILD", + sha256 = CUDA_TILE_SHA256, + strip_prefix = "cuda-tile-{}".format(CUDA_TILE_COMMIT), + urls = tf_mirror_urls("https://github.com/NVIDIA/cuda-tile/archive/{}.tar.gz".format(CUDA_TILE_COMMIT)), + ) diff --git a/third_party/xla/third_party/tensor_ir/BUILD.bazel b/third_party/xla/third_party/tensor_ir/BUILD.bazel new file mode 100644 index 00000000000000..e69de29bb2d1d6 diff --git a/third_party/xla/third_party/tensor_ir/tensor_ir.BUILD b/third_party/xla/third_party/tensor_ir/tensor_ir.BUILD new file mode 100644 index 00000000000000..3d591efd26df72 --- /dev/null +++ b/third_party/xla/third_party/tensor_ir/tensor_ir.BUILD @@ -0,0 +1,480 @@ +load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library") +load("@rules_cc//cc:cc_library.bzl", "cc_library") + +exports_files(["LICENSE"]) + +td_library( + name = "TensorIRTdFiles", + srcs = glob(["include/tensor_ir/Dialect/*.td"]), + includes = ["include"], + deps = [ + "@llvm-project//mlir:BuiltinDialectTdFiles", + "@llvm-project//mlir:ControlFlowInterfacesTdFiles", + "@llvm-project//mlir:FunctionInterfacesTdFiles", + "@llvm-project//mlir:GPUOpsTdFiles", + "@llvm-project//mlir:InferTypeOpInterfaceTdFiles", + "@llvm-project//mlir:OpBaseTdFiles", + "@llvm-project//mlir:SideEffectInterfacesTdFiles", + ], +) + +gentbl_cc_library( + name = "TensorIRDialectIncGen", + tbl_outs = [ + ( + [ + "-gen-dialect-decls", + "-dialect=nv_tensor_ir", + ], + "include/tensor_ir/Dialect/TensorDialect.h.inc", + ), + ( + [ + "-gen-dialect-defs", + "-dialect=nv_tensor_ir", + ], + "include/tensor_ir/Dialect/TensorDialect.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/tensor_ir/Dialect/TensorDialect.td", + deps = [":TensorIRTdFiles"], +) + +gentbl_cc_library( + name = "TensorIROpsIncGen", + tbl_outs = [ + ( + ["-gen-op-decls"], + "include/tensor_ir/Dialect/TensorOps.h.inc", + ), + ( + ["-gen-op-defs"], + "include/tensor_ir/Dialect/TensorOps.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/tensor_ir/Dialect/TensorOps.td", + deps = [":TensorIRTdFiles"], +) + +gentbl_cc_library( + name = "TensorIREnumsIncGen", + tbl_outs = [ + ( + ["-gen-enum-decls"], + "include/tensor_ir/Dialect/TensorEnums.h.inc", + ), + ( + ["-gen-enum-defs"], + "include/tensor_ir/Dialect/TensorEnums.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/tensor_ir/Dialect/TensorEnums.td", + deps = [":TensorIRTdFiles"], +) + +gentbl_cc_library( + name = "TensorIROpInterfacesIncGen", + tbl_outs = [ + ( + ["-gen-op-interface-decls"], + "include/tensor_ir/Dialect/TensorOpInterfaces.h.inc", + ), + ( + ["-gen-op-interface-defs"], + "include/tensor_ir/Dialect/TensorOpInterfaces.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/tensor_ir/Dialect/TensorInterfaces.td", + deps = [":TensorIRTdFiles"], +) + +gentbl_cc_library( + name = "TensorIRAttrInterfacesIncGen", + tbl_outs = [ + ( + ["-gen-attr-interface-decls"], + "include/tensor_ir/Dialect/TensorAttrInterfaces.h.inc", + ), + ( + ["-gen-attr-interface-defs"], + "include/tensor_ir/Dialect/TensorAttrInterfaces.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/tensor_ir/Dialect/TensorAttrInterfaces.td", + deps = [":TensorIRTdFiles"], +) + +gentbl_cc_library( + name = "TensorIRAttrsIncGen", + tbl_outs = [ + ( + [ + "-gen-attrdef-decls", + "-attrdefs-dialect=nv_tensor_ir", + ], + "include/tensor_ir/Dialect/TensorAttrs.h.inc", + ), + ( + [ + "-gen-attrdef-defs", + "-attrdefs-dialect=nv_tensor_ir", + ], + "include/tensor_ir/Dialect/TensorAttrs.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/tensor_ir/Dialect/TensorAttrs.td", + deps = [":TensorIRTdFiles"], +) + +gentbl_cc_library( + name = "TensorIROpsCanonicalizationIncGen", + strip_include_prefix = "lib/Dialect", + tbl_outs = [ + ( + ["-gen-rewriters"], + "lib/Dialect/TensorOpsCanonicalization.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "lib/Dialect/TensorOpsCanonicalization.td", + deps = [ + ":TensorIRTdFiles", + "@llvm-project//mlir:OpBaseTdFiles", + ], +) + +gentbl_cc_library( + name = "TensorIRTransformPassesIncGen", + tbl_outs = [ + ( + [ + "-gen-pass-decls", + "-name=NVTensorIRTransform", + ], + "include/tensor_ir/Transform/Passes.h.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/tensor_ir/Transform/Passes.td", + deps = [ + ":TensorIRTdFiles", + "@llvm-project//mlir:PassBaseTdFiles", + ], +) + +gentbl_cc_library( + name = "TensorToCudaTileConversionPassIncGen", + tbl_outs = [ + ( + [ + "-gen-pass-decls", + "-name=TensorToCudaTileConversion", + ], + "include/tensor_ir/Conversion/TensorToCudaTile/Passes.h.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "include/tensor_ir/Conversion/TensorToCudaTile/Passes.td", + deps = [ + ":TensorIRTdFiles", + "@llvm-project//mlir:PassBaseTdFiles", + ], +) + +cc_library( + name = "NVTensorIRSupport", + srcs = ["lib/Support/TCutegen.cpp"], + hdrs = [ + "include/tensor_ir/Support/Macros.h", + "include/tensor_ir/Support/Status.h", + "include/tensor_ir/Support/TCutegen.h", + ], + includes = ["include"], + visibility = ["//visibility:public"], + deps = [ + "@llvm-project//llvm:Support", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:Support", + ], +) + +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 = [ + "lib/Dialect/Canonicalization.cpp", + "lib/Dialect/TensorAttrs.cpp", + "lib/Dialect/TensorDialect.cpp", + "lib/Dialect/TensorOps.cpp", + ], + hdrs = [ + "include/tensor_ir/Dialect/TensorIR.h", + "include/tensor_ir/Dialect/TensorIRAttrs.h", + ], + includes = ["include"], + visibility = ["//visibility:public"], + deps = [ + ":NVTensorIRSupport", + ":TensorIRAttrInterfacesIncGen", + ":TensorIRAttrsIncGen", + ":TensorIRDialectIncGen", + ":TensorIREnumsIncGen", + ":TensorIROpInterfacesIncGen", + ":TensorIROpsCanonicalizationIncGen", + ":TensorIROpsIncGen", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:BytecodeOpInterface", + "@llvm-project//mlir:ControlFlowInterfaces", + "@llvm-project//mlir:FunctionInterfaces", + "@llvm-project//mlir:GPUDialect", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:InferTypeOpInterface", + "@llvm-project//mlir:SideEffectInterfaces", + ], +) + +cc_library( + name = "NVTensorIRUtils", + srcs = [ + "lib/Utils/ComputeCapability.cpp", + "lib/Utils/Utils.cpp", + ], + hdrs = [ + "include/tensor_ir/Utils/ComputeCapability.h", + "include/tensor_ir/Utils/Utils.h", + ], + includes = ["include"], + visibility = ["//visibility:public"], + deps = [ + ":NVTensorIRDialect", + ":NVTensorIRSupport", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:ArithDialect", + "@llvm-project//mlir:BytecodeWriter", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:Support", + ], +) + +cc_library( + name = "NVTensorIRAnalysis", + srcs = [ + "lib/Analysis/KernelArgLayout.cpp", + "lib/Analysis/TileAnalyzer.cpp", + "lib/Analysis/TileCandidateGenerator.cpp", + ], + hdrs = [ + "include/tensor_ir/Analysis/TileAnalyzer.h", + "include/tensor_ir/Analysis/TileCandidateGenerator.h", + "include/tensor_ir/Compiler/CudaTile/KernelArgLayout.h", + ], + includes = ["include"], + visibility = ["//visibility:public"], + deps = [ + ":NVTensorIRDialect", + ":NVTensorIRRuntime", + ":NVTensorIRUtils", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:IR", + ], +) + +cc_library( + name = "NVTensorIRTransform", + srcs = glob(["lib/Transform/*.cpp"]), + hdrs = ["include/tensor_ir/Transform/Passes.h"], + includes = ["include"], + visibility = ["//visibility:public"], + deps = [ + ":NVTensorIRAnalysis", + ":NVTensorIRDialect", + ":NVTensorIRSupport", + ":NVTensorIRUtils", + ":TensorIRTransformPassesIncGen", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:FuncDialect", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:Support", + ], +) + +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", + ], + includes = ["include"], + visibility = ["//visibility:public"], + deps = [ + ":NVTensorIRAnalysis", + ":NVTensorIRDialect", + ":NVTensorIRSupport", + ":NVTensorIRUtils", + ":TensorToCudaTileConversionPassIncGen", + "@cuda_tile//:CudaTileDialect", + "@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"], + includes = ["include"], + visibility = ["//visibility:public"], + deps = [ + ":NVTensorIRDialect", + ":NVTensorIRToCudaTileConversion", + ":NVTensorIRTransform", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:Transforms", + ], +) + +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", + ], + 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"]), + 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", + ], + visibility = ["//visibility:public"], + deps = [ + ":NVTensorIRCudaApi", + ":NVTensorIRDialect", + ":NVTensorIRRuntime", + ":NVTensorIRSupport", + ":NVTensorIRUtils", + "@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:FuncDialect", + "@llvm-project//mlir:FuncExtensions", + "@llvm-project//mlir:IR", + ], +) + +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", + ], + includes = ["include"], + visibility = ["//visibility:public"], + deps = [ + ":NVTensorIRAnalysis", + ":NVTensorIRCudaTilePipelines", + ":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", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:Support", + ], +) + +cc_library( + name = "NVTensorIRCAPI", + srcs = ["lib/CAPI/TensorIR.cpp"], + hdrs = ["include/tensor_ir-c/TensorIR.h"], + includes = ["include"], + visibility = ["//visibility:public"], + deps = [ + ":NVTensorIRCompiler", + ":NVTensorIRDialect", + ":NVTensorIRRegistration", + ":NVTensorIRRuntime", + "@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 new file mode 100644 index 00000000000000..089aacd7ceafe8 --- /dev/null +++ b/third_party/xla/third_party/tensor_ir/workspace.bzl @@ -0,0 +1,16 @@ +"""Provides the repository macro to import Tensor IR.""" + +load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls") + +def repo(): + """Imports Tensor IR.""" + TENSOR_IR_COMMIT = "63692d79629e6f32a1d8757695590a59e0adbafd" + TENSOR_IR_SHA256 = "b80794d7c2bfb1bc1ca432d892977becba8d35ec6c18c586acbd648ccc8074dd" + + tf_http_archive( + name = "tensor_ir", + build_file = "//third_party/tensor_ir:tensor_ir.BUILD", + sha256 = TENSOR_IR_SHA256, + strip_prefix = "tensor-ir-{}".format(TENSOR_IR_COMMIT), + urls = tf_mirror_urls("https://github.com/NVIDIA/tensor-ir/archive/{}.tar.gz".format(TENSOR_IR_COMMIT)), + ) diff --git a/third_party/xla/workspace2.bzl b/third_party/xla/workspace2.bzl index 923de0b6bda95d..0f49a1db57c284 100644 --- a/third_party/xla/workspace2.bzl +++ b/third_party/xla/workspace2.bzl @@ -31,6 +31,7 @@ load("//third_party/brotli:workspace.bzl", brotli = "repo") load("//third_party/clang_toolchain:cc_configure_clang.bzl", "cc_download_clang_toolchain") load("//third_party/compute_library:workspace.bzl", compute_library = "repo") load("//third_party/cpuinfo:workspace.bzl", cpuinfo = "repo") +load("//third_party/cuda_tile:workspace.bzl", cuda_tile = "repo") load("//third_party/cudnn_frontend:workspace.bzl", cudnn_frontend = "repo") load("//third_party/cutlass:workspace.bzl", cutlass = "repo") load("//third_party/cutlass_cutedsl_runtime:workspace.bzl", cutlass_cutedsl_runtime = "repo") @@ -81,6 +82,7 @@ load("//third_party/spdlog:workspace.bzl", spdlog = "repo") load("//third_party/sqlite:workspace.bzl", sqlite = "repo") load("//third_party/stablehlo:workspace.bzl", stablehlo = "repo") load("//third_party/system_libpci:workspace.bzl", system_libpci = "repo") +load("//third_party/tensor_ir:workspace.bzl", tensor_ir = "repo") load("//third_party/tensorrt:tensorrt_configure.bzl", "tensorrt_configure") load("//third_party/tensorrt:workspace.bzl", tensorrt = "repo") load("//third_party/transformer_engine:workspace.bzl", transformer_engine = "repo") @@ -147,11 +149,13 @@ def _initialize_third_party(): spdlog() sqlite() stablehlo() + tensor_ir() tensorrt() transformer_engine() triton() uv() xnnpack() + cuda_tile() cutlass() cutlass_cutedsl_runtime() From 59ab457e356cdb82f57d4eaebf6c9121ff13d965 Mon Sep 17 00:00:00 2001 From: Dirk Hornung Date: Thu, 27 Aug 2026 05:12:33 -0700 Subject: [PATCH 32/32] [XLA:GPU] Disable epilogue convolution fusions on pre-Ampere devices. Note: Alternative would be to add input channel padding as FP16 Tensor Core execution plans require channel to be 128bit aligned. PiperOrigin-RevId: 971890275 --- .../transforms/conv_fusion_rewriter_test.cc | 36 +++++++++++++++++++ .../gpu/transforms/cudnn_fusion_utils.cc | 8 +++++ 2 files changed, 44 insertions(+) diff --git a/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter_test.cc b/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter_test.cc index 91baddc0f6ab36..59a515813eb372 100644 --- a/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter_test.cc +++ b/third_party/xla/xla/backends/gpu/transforms/conv_fusion_rewriter_test.cc @@ -2416,6 +2416,42 @@ TEST_F(ConvFusionRewriterIntegrationTest, )"); } +TEST_F(ConvFusionRewriterUnitTest, EpilogueNotFusedOnPreAmpere) { + const char* const hlo_string = R"( + HloModule Test + + ENTRY Test { + input = f16[1,16,16,16] parameter(0) + filter = f16[3,3,16,32] parameter(1) + conv = f16[1,16,16,32] convolution(input, filter), + window={size=3x3 pad=1_1x1_1}, + dim_labels=b01f_01io->b01f + bias = f16[32] parameter(2) + bcast = f16[1,16,16,32] broadcast(bias), dimensions={3} + ROOT add = f16[1,16,16,32] add(conv, bcast) + })"; + + se::DeviceDescription volta_device; + volta_device.set_gpu_compute_capability(se::CudaComputeCapability::Volta()); + + // On Volta (SM70), epilogue should NOT be fused into the custom fusion. + RunAndMatch(hlo_string, + m::Add(m::Fusion(m::Parameter(0), m::Parameter(1)) + .WithFusionKind(HloInstruction::FusionKind::kCustom), + m::Broadcast(m::Parameter(2))), + /*run_algebraic_simplifier=*/false, volta_device); + + se::DeviceDescription ampere_device; + ampere_device.set_gpu_compute_capability(se::CudaComputeCapability::Ampere()); + + // On Ampere+ (SM80+), epilogue SHOULD be fused into the custom fusion. + RunAndMatch(hlo_string, + m::Fusion(m::Parameter(0), m::Parameter(1), m::Parameter(2)) + .WithFusionKind(HloInstruction::FusionKind::kCustom) + .WithShape(F16, {1, 16, 16, 32}), + /*run_algebraic_simplifier=*/false, ampere_device); +} + } // namespace } // namespace gpu } // namespace xla diff --git a/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_utils.cc b/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_utils.cc index 3b5495b800ed89..a60ab8e6ad60fb 100644 --- a/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_utils.cc +++ b/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_utils.cc @@ -87,6 +87,14 @@ bool IsEpilogueOpSupportedByCuDNN(const HloInstruction& hlo, if (is_nchw) { return false; } + + // cuDNN only supports epilogues on Ampere and above. + const se::CudaComputeCapability* cuda_cc = + device_info.gpu_compute_capability().cuda_compute_capability(); + if (cuda_cc != nullptr && !cuda_cc->IsAtLeastAmpere()) { + return false; + } + const HloOpcode opcode = hlo.opcode(); // Do not fuse chained converts (a convert whose operand is already a // convert). Note: Fusing chained converts could steal a convert from an