Skip to content

[pull] master from tensorflow:master #8838

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
pull merged 13 commits into from
Sep 19, 2026
Merged

[pull] master from tensorflow:master #8838

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
ea17f4c
[Mosaic] Add helpers to infer sequential accesses from tiled layouts
naummo Sep 18, 2026
2e368db
[XLA:GPU] Fix scatter with empty index vector in ScatterWithDistribut…
akuegel Sep 18, 2026
639345e
Remove exclusive-if-local test tags to eliminate test serialization b…
dmiltr3 Sep 18, 2026
89fefb6
[tf.data] Propagate context when creating tf.data threads.
yangustc07 Sep 18, 2026
a962b25
Increase MSan memory limits and test sharding to resolve OOMs and tim…
dmiltr3 Sep 18, 2026
b53d234
Miscellaneous fixes for v5p and v6e.
tensorflower-gardener Sep 18, 2026
1656f20
Replace `/` with `_` in generated `lit_test` target names
mrguenther Sep 18, 2026
07e2326
Fix the ping-pong bug in OptimizeBatchMatmulPass
majiddadashi Sep 18, 2026
2e6f1ef
Increase shard_count of np_indexing_test to 30 to fix TSAN timeouts.
tensorflower-gardener Sep 18, 2026
720136c
Round the device memory size to the nearest GiB while producing topol…
junwhanahn Sep 18, 2026
10b8781
Move USE_MATH_DEFINES macro above the includes on hlo evaluator.
ekayaaslan Sep 18, 2026
2ab8f03
Before CheckLayouts(), update computation_layouts_ with latest entry …
tensorflower-gardener Sep 19, 2026
62ffa09
Copy removal: count a view's readers as uses of the viewed value
patrick-toulme Sep 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 33 additions & 0 deletions tensorflow/compiler/mlir/lite/tests/optimize_batch_matmul.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -242,4 +242,37 @@ func.func @FuseBatchMatmulToTransposeNegative(%arg0: tensor<2048x32x1x8x2xf32>,
%196 = "tfl.batch_matmul"(%arg1, %195) <{adj_x = false, adj_y = false, asymmetric_quantize_inputs = false}> : (tensor<2048x1x32x2xf32>, tensor<2048x1x2x256xf32>) -> tensor<2048x1x32x256xf32>
return %196 : tensor<2048x1x32x256xf32>
// CHECK: "tfl.transpose"
}

// CHECK-LABEL: NotFuseTransposeFCRhsToBatchMatmulResourceConstant
func.func @NotFuseTransposeFCRhsToBatchMatmulResourceConstant(%arg0: tensor<16x1024xf32>, %arg1: none) -> tensor<16x128xf32> {
%cst = arith.constant dense<[1, 0]> : tensor<2xi32>
%weight = arith.constant dense_resource<__elided__> : tensor<1024x128xf32>
%0 = "tfl.transpose"(%weight, %cst) : (tensor<1024x128xf32>, tensor<2xi32>) -> tensor<128x1024xf32>
// CHECK: "tfl.fully_connected"
// CHECK-NOT: "tfl.batch_matmul"
%1 = "tfl.fully_connected"(%arg0, %0, %arg1) {asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<16x1024xf32>, tensor<128x1024xf32>, none) -> tensor<16x128xf32>
func.return %1 : tensor<16x128xf32>
}

// CHECK-LABEL: NotFuseTransposeFCRhsToBatchMatmulReshapedResourceConstant
func.func @NotFuseTransposeFCRhsToBatchMatmulReshapedResourceConstant(%arg0: tensor<16x1024xf32>, %arg1: none) -> tensor<16x128xf32> {
%cst = arith.constant dense<[1, 0]> : tensor<2xi32>
%cst_shape = arith.constant dense<[1024, 128]> : tensor<2xi32>
%weight = arith.constant dense_resource<__elided__> : tensor<1024x128xf32>
%0 = "tfl.reshape"(%weight, %cst_shape) : (tensor<1024x128xf32>, tensor<2xi32>) -> tensor<1024x128xf32>
%1 = "tfl.transpose"(%0, %cst) : (tensor<1024x128xf32>, tensor<2xi32>) -> tensor<128x1024xf32>
// CHECK: "tfl.fully_connected"
// CHECK-NOT: "tfl.batch_matmul"
%2 = "tfl.fully_connected"(%arg0, %1, %arg1) {asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<16x1024xf32>, tensor<128x1024xf32>, none) -> tensor<16x128xf32>
func.return %2 : tensor<16x128xf32>
}

// CHECK-LABEL: BatchMatmulResourceConstantToFullyConnected
func.func @BatchMatmulResourceConstantToFullyConnected(%arg0: tensor<16x1024xf32>) -> tensor<16x128xf32> {
%weight = arith.constant dense_resource<__elided__> : tensor<1024x128xf32>
// CHECK: "tfl.fully_connected"
// CHECK-NOT: "tfl.batch_matmul"
%0 = "tfl.batch_matmul"(%arg0, %weight) {adj_x = false, adj_y = false, asymmetric_quantize_inputs = false} : (tensor<16x1024xf32>, tensor<1024x128xf32>) -> tensor<16x128xf32>
func.return %0 : tensor<16x128xf32>
}
9 changes: 7 additions & 2 deletions tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ limitations under the License.
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
#include "mlir/IR/Matchers.h" // from @llvm-project
#include "mlir/IR/OpDefinition.h" // from @llvm-project
#include "mlir/IR/PatternMatch.h" // from @llvm-project
#include "mlir/IR/TypeUtilities.h" // from @llvm-project
#include "mlir/IR/Value.h" // from @llvm-project
Expand All @@ -50,7 +51,9 @@ bool NotFromFoldableChain(mlir::Value value) {
mlir::Operation* defining_op = value.getDefiningOp();

while (defining_op) {
if (mlir::isa<DequantizeOp>(defining_op)) {
if (mlir::isa<DequantizeOp>(defining_op) ||
defining_op->hasTrait<mlir::OpTrait::ConstantLike>() ||
matchPattern(defining_op, mlir::m_Constant())) {
return false;
}

Expand All @@ -59,8 +62,10 @@ bool NotFromFoldableChain(mlir::Value value) {
defining_op = reshape_op.getInput().getDefiningOp();
} else if (auto split_op = mlir::dyn_cast<SplitOp>(defining_op)) {
defining_op = split_op.getValue().getDefiningOp();
} else if (auto transpose_op = mlir::dyn_cast<TransposeOp>(defining_op)) {
defining_op = transpose_op.getInput().getDefiningOp();
} else {
// Stop if the op is not Dequantize, Reshape, or Split.
// Stop if the op does not preserve constant nature.
break;
}
}
Expand Down
3 changes: 0 additions & 3 deletions tensorflow/core/kernels/data/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -811,7 +811,6 @@ tf_kernel_library(
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/data:dataset_utils",
"//tensorflow/core/data:name_utils",
"//tensorflow/core/data:stats_utils",
Expand Down Expand Up @@ -968,11 +967,9 @@ tf_kernel_library(
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/data:captured_function",
"//tensorflow/core/data:dataset_utils",
"//tensorflow/core/data:name_utils",
"//tensorflow/core/data:stats_utils",
"//tensorflow/core/framework:attr_value_proto_cc",
"//tensorflow/core/framework:types_proto_cc",
"//tensorflow/core/platform:env",
Expand Down
4 changes: 3 additions & 1 deletion tensorflow/core/kernels/data/experimental/map_and_batch_dataset_op.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ limitations under the License.
#include "tensorflow/core/platform/stringprintf.h"
#include "tensorflow/core/profiler/lib/traceme.h"
#include "tensorflow/core/profiler/lib/traceme_encode.h"
#include "tsl/platform/context.h"

namespace tensorflow {
namespace data {
Expand Down Expand Up @@ -528,7 +529,8 @@ class MapAndBatchDatasetOp::Dataset : public DatasetBase {
auto new_ctx = std::make_shared<IteratorContext>(*ctx);
runner_thread_.reset(Env::Default()->StartThread(
/*thread_options=*/{}, kTFDataMapAndBatch,
std::bind(&Iterator::RunnerThread, this, new_ctx)));
tsl::WithCurrentContext(
std::bind(&Iterator::RunnerThread, this, new_ctx))));
}
}

Expand Down
7 changes: 5 additions & 2 deletions tensorflow/core/kernels/data/experimental/parallel_interleave_dataset_op.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ limitations under the License.
#include "tensorflow/core/platform/stringprintf.h"
#include "tensorflow/core/profiler/lib/traceme.h"
#include "tensorflow/core/profiler/lib/traceme_encode.h"
#include "tsl/platform/context.h"

namespace tensorflow {
namespace data {
Expand Down Expand Up @@ -574,7 +575,8 @@ class ParallelInterleaveDatasetOp::Dataset : public DatasetBase {
worker_threads_.emplace_back(Env::Default()->StartThread(
/*thread_options=*/{},
absl::StrCat(kDataParallelInterleaveWorker, "_", i),
[this, new_ctx, i]() { WorkerThread(new_ctx, i); }));
tsl::WithCurrentContext(
[this, new_ctx, i]() { WorkerThread(new_ctx, i); })));
}
}
return absl::OkStatus();
Expand Down Expand Up @@ -697,7 +699,8 @@ class ParallelInterleaveDatasetOp::Dataset : public DatasetBase {
worker_threads_.emplace_back(Env::Default()->StartThread(
/*thread_options=*/{},
absl::StrCat(kDataParallelInterleaveWorker, "_", i),
[this, new_ctx, i]() { WorkerThread(new_ctx, i); }));
tsl::WithCurrentContext(
[this, new_ctx, i]() { WorkerThread(new_ctx, i); })));
}
DCHECK(interleave_indices_.size() == dataset()->cycle_length_);
DCHECK(staging_indices_.size() == dataset()->prefetch_input_elements_);
Expand Down
7 changes: 5 additions & 2 deletions tensorflow/core/kernels/data/experimental/parse_example_dataset_op.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ limitations under the License.
#include "tensorflow/core/profiler/lib/traceme.h"
#include "tensorflow/core/profiler/lib/traceme_encode.h"
#include "tensorflow/core/util/example_proto_fast_parsing.h"
#include "tsl/platform/context.h"

namespace tensorflow {
namespace data {
Expand Down Expand Up @@ -571,11 +572,13 @@ class ParseExampleDatasetOp : public UnaryDatasetOpKernel {
auto ctx_copy = std::make_shared<IteratorContext>(*ctx);
runner_thread_.reset(Env::Default()->StartThread(
/*thread_options=*/{}, "tf_data_parallel_map",
std::bind(&Iterator::RunnerThread, this, ctx_copy)));
tsl::WithCurrentContext(
std::bind(&Iterator::RunnerThread, this, ctx_copy))));
if (ctx->stats_aggregator()) {
stats_thread_.reset(Env::Default()->StartThread(
/*thread_options=*/{}, "tf_data_parallel_map_stats",
std::bind(&Iterator::StatsThread, this, ctx_copy)));
tsl::WithCurrentContext(
std::bind(&Iterator::StatsThread, this, ctx_copy))));
}
}
}
Expand Down
4 changes: 3 additions & 1 deletion tensorflow/core/kernels/data/parallel_batch_dataset_op.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@ limitations under the License.
#include "tensorflow/core/profiler/lib/traceme.h"
#include "tensorflow/core/profiler/lib/traceme_encode.h"
#include "tensorflow/core/util/batch_util.h"
#include "tsl/platform/context.h"

namespace tensorflow {
namespace data {
Expand Down Expand Up @@ -480,7 +481,8 @@ class ParallelBatchDatasetOp::Dataset : public DatasetBase {
auto new_ctx = std::make_shared<IteratorContext>(*ctx);
runner_thread_.reset(Env::Default()->StartThread(
/*thread_options=*/{}, kTFDataParallelBatch,
std::bind(&Iterator::RunnerThread, this, new_ctx)));
tsl::WithCurrentContext(
std::bind(&Iterator::RunnerThread, this, new_ctx))));
}
}

Expand Down
4 changes: 3 additions & 1 deletion tensorflow/core/kernels/data/parallel_filter_dataset_op.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ limitations under the License.
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/profiler/lib/traceme.h"
#include "tensorflow/core/profiler/lib/traceme_encode.h"
#include "tsl/platform/context.h"

namespace tensorflow {
namespace data {
Expand Down Expand Up @@ -336,7 +337,8 @@ class ParallelFilterDatasetOp::Dataset : public DatasetBase {
auto ctx_copy = std::make_shared<IteratorContext>(*ctx);
runner_thread_.reset(Env::Default()->StartThread(
/*thread_options=*/{}, "tf_data_parallel_filter",
std::bind(&Iterator::RunnerThread, this, ctx_copy)));
tsl::WithCurrentContext(
std::bind(&Iterator::RunnerThread, this, ctx_copy))));
}
}

Expand Down
7 changes: 5 additions & 2 deletions tensorflow/core/kernels/data/parallel_map_dataset_op.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ limitations under the License.
#include "tensorflow/core/profiler/lib/traceme.h"
#include "tensorflow/core/profiler/lib/traceme_encode.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"
#include "tsl/platform/context.h"

namespace tensorflow {
namespace data {
Expand Down Expand Up @@ -538,11 +539,13 @@ class ParallelMapDatasetOp::Dataset : public DatasetBase {
auto ctx_copy = std::make_shared<IteratorContext>(*ctx);
runner_thread_.reset(Env::Default()->StartThread(
/*thread_options=*/{}, "tf_data_parallel_map",
std::bind(&Iterator::RunnerThread, this, ctx_copy)));
tsl::WithCurrentContext(
std::bind(&Iterator::RunnerThread, this, ctx_copy))));
if (ctx->stats_aggregator()) {
stats_thread_.reset(Env::Default()->StartThread(
/*thread_options=*/{}, "tf_data_parallel_map_stats",
std::bind(&Iterator::StatsThread, this, ctx_copy)));
tsl::WithCurrentContext(
std::bind(&Iterator::StatsThread, this, ctx_copy))));
}
}
}
Expand Down
4 changes: 3 additions & 1 deletion tensorflow/core/kernels/data/prefetch_dataset_op.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@ limitations under the License.
#include "tensorflow/core/profiler/lib/traceme.h"
#include "tensorflow/core/profiler/lib/traceme_encode.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"
#include "tsl/platform/context.h"
#include "tsl/platform/mutex.h"

namespace tensorflow {
Expand Down Expand Up @@ -544,7 +545,8 @@ class PrefetchDatasetOp::Dataset : public DatasetBase {
std::make_shared<IteratorContext>(*ctx);
prefetch_thread_.reset(Env::Default()->StartThread(
/*thread_options=*/{}, "tf_data_prefetch",
[this, new_ctx]() { PrefetchThread(new_ctx); }));
tsl::WithCurrentContext(
[this, new_ctx]() { PrefetchThread(new_ctx); })));
}
return absl::OkStatus();
}
Expand Down
1 change: 1 addition & 0 deletions tensorflow/python/distribute/integration_test/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ distribute_py_strict_test(
exec_properties = select({
# copybara:uncomment_begin(google-only)
# "@bazel_tools//tools/cpp:asan_build": {"mem": "16g"},
# "//tools/cpp:msan_build": {"mem": "20g"},
# copybara:uncomment_end
"//conditions:default": None,
}),
Expand Down
6 changes: 6 additions & 0 deletions tensorflow/python/eager/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1148,6 +1148,12 @@ cuda_py_strict_test(
name = "remote_test",
size = "medium",
srcs = ["remote_test.py"],
exec_properties = select({
# copybara:uncomment_begin(google-only)
# "//tools/cpp:msan_build": {"mem": "30g"},
# copybara:uncomment_end
"//conditions:default": None,
}),
grpc_enabled = True,
shard_count = 2,
tags = [
Expand Down
3 changes: 0 additions & 3 deletions tensorflow/python/kernel_tests/distributions/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -326,9 +326,6 @@ cuda_py_strict_test(
"XLA_PYTHON_CLIENT_PREALLOCATE": "false",
},
shard_count = 5,
tags = [
"exclusive-if-local",
],
xla_enable_strict_auto_jit = False,
deps = [
"//tensorflow/python/eager:backprop",
Expand Down
2 changes: 1 addition & 1 deletion tensorflow/python/kernel_tests/linalg/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -202,7 +202,7 @@ cuda_py_strict_test(
name = "linear_operator_block_diag_test",
size = "medium",
srcs = ["linear_operator_block_diag_test.py"],
shard_count = 8,
shard_count = 16,
tags = ["optonly"],
deps = [
"//tensorflow/python/eager:backprop",
Expand Down
3 changes: 0 additions & 3 deletions tensorflow/python/kernel_tests/nn_ops/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -909,9 +909,6 @@ cuda_py_strict_test(
"TF_FORCE_GPU_ALLOW_GROWTH": "true",
"XLA_PYTHON_CLIENT_PREALLOCATE": "false",
},
tags = [
"exclusive-if-local",
],
deps = [
":xent_op_test_base",
"//tensorflow/python/client:session",
Expand Down
3 changes: 1 addition & 2 deletions tensorflow/python/ops/numpy_ops/tests/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -222,7 +222,6 @@ py_test(
shard_count = 20,
strict_deps = True,
tags = [
"exclusive-if-local",
"gpu",
"no_pip",
"requires-gpu-nvidia",
Expand Down Expand Up @@ -253,7 +252,7 @@ py_test(
"--num_generated_cases=90",
"--enable_x64", # Needed to enable dtype check
],
shard_count = 10,
shard_count = 30,
strict_deps = True,
# TODO(b/164245103): Re-enable GPU once tf.tensor_strided_slice_update's segfault is fixed.
tags = [
Expand Down
Loading
Loading