From a4b6c5623bc1e8b232a8c301f3b35da33b6afea8 Mon Sep 17 00:00:00 2001 From: Hyeontaek Lim Date: Fri, 28 Aug 2026 16:49:17 -0700 Subject: [PATCH 01/12] [IFRT] Use canonical custom layouts on array_impl tests The custom layouts used in array_impl/xla_array_impl were not fully canonicalized and was relying on the runtime's implicit layout canonicalization. This change uses a tiled layout for the custom layout tests when running on TPU so that the runtime can apply strong validation for canonicalized layouts and does not have to perform implicit layout canonicalization. PiperOrigin-RevId: 972856960 --- .../xla/python/ifrt/array_impl_test_lib.cc | 33 ++++++++----------- .../pjrt_ifrt/xla_array_impl_test_lib.cc | 32 ++++++++++++++---- 2 files changed, 38 insertions(+), 27 deletions(-) diff --git a/third_party/xla/xla/python/ifrt/array_impl_test_lib.cc b/third_party/xla/xla/python/ifrt/array_impl_test_lib.cc index 641f51ba891af8..fb30db8b988f90 100644 --- a/third_party/xla/xla/python/ifrt/array_impl_test_lib.cc +++ b/third_party/xla/xla/python/ifrt/array_impl_test_lib.cc @@ -34,7 +34,7 @@ limitations under the License. #include "absl/time/clock.h" #include "absl/time/time.h" #include "absl/types/span.h" -#include "xla/layout_util.h" +#include "xla/layout.h" #include "xla/pjrt/pjrt_layout.h" #include "xla/python/ifrt/array.h" #include "xla/python/ifrt/array_spec.h" @@ -954,27 +954,20 @@ TEST(ArrayImplTest, MakeArraysFromHostBufferShardsWithLayout) { absl::c_iota(data, 0); Device* device = client->addressable_devices()[0]; - std::shared_ptr layout; + std::shared_ptr layout; int64_t expected_size; { - auto xla_layout = - xla::LayoutUtil::MakeDescendingLayout(shape.dims().size()); - TF_ASSERT_OK_AND_ASSIGN(auto device_list, client->MakeDeviceList({device})); - auto topology = client->GetTopologyForDevices(device_list); - if (topology.ok()) { - auto topology_desc = (*topology)->description(); - TF_ASSERT_OK_AND_ASSIGN( - xla::Shape xla_shape, - topology_desc->MakeCanonicalShapeForMemorySpace( - topology_desc->GetDefaultMemorySpaceKindId(), - xla::ShapeUtil::MakeShape(xla::PrimitiveType::F32, shape.dims()), - &xla_layout)); - layout = std::make_shared(xla_shape.layout()); - expected_size = xla::ShapeUtil::ArraySize(xla_shape); - } else { - layout = std::make_shared(std::move(xla_layout)); - expected_size = *dtype.byte_size() * shape.num_elements(); - } + ASSERT_OK_AND_ASSIGN(std::shared_ptr default_layout, + client->GetDefaultPjRtLayout(dtype, shape.dims(), + device, MemoryKind())); + xla::Shape xla_shape = + xla::ShapeUtil::MakeShape(xla::PrimitiveType::F32, shape.dims()); + *xla_shape.mutable_layout() = default_layout->xla_layout(); + // We assume that reversing the minor_to_major still gives a valid layout + // for this shape. + absl::c_reverse(*xla_shape.mutable_layout()->mutable_minor_to_major()); + layout = std::make_shared(xla_shape.layout()); + expected_size = xla::ShapeUtil::ArraySize(xla_shape); } ArrayRef array; diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_array_impl_test_lib.cc b/third_party/xla/xla/python/pjrt_ifrt/xla_array_impl_test_lib.cc index 9919a1b68d84db..3447807509b7c1 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/xla_array_impl_test_lib.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/xla_array_impl_test_lib.cc @@ -16,6 +16,7 @@ limitations under the License. #include #include #include +#include #include #include @@ -24,7 +25,7 @@ limitations under the License. #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/types/span.h" -#include "xla/layout_util.h" +#include "xla/layout.h" #include "xla/pjrt/pjrt_compiler.h" #include "xla/pjrt/pjrt_layout.h" #include "xla/python/ifrt/array.h" @@ -62,9 +63,15 @@ TEST_P(XlaArrayImplHashTest, HashValuesDifferentLayouts) { absl::c_iota(data, 0.0f); // Array 0: Row-major layout - std::shared_ptr layout0 = - xla::ifrt::PjRtLayout::Create(std::make_shared( - xla::LayoutUtil::MakeDescendingLayout(/*num_dims=*/2))); + std::shared_ptr layout0; + { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr default_layout, + client->GetDefaultPjRtLayout(dtype, shape.dims(), + client->addressable_devices().at(0), + MemoryKind())); + layout0 = xla::ifrt::PjRtLayout::Create(default_layout); + } ASSERT_OK_AND_ASSIGN( ArrayRef array0, client->MakeArrayFromHostBuffer( @@ -83,9 +90,20 @@ TEST_P(XlaArrayImplHashTest, HashValuesDifferentLayouts) { ASSERT_OK(actual_layout0.status()); // Array 1: Column-major layout - std::shared_ptr layout1 = - xla::ifrt::PjRtLayout::Create(std::make_shared( - xla::LayoutUtil::MakeAscendingLayout(/*num_dims=*/2))); + std::shared_ptr layout1; + { + ASSERT_OK_AND_ASSIGN( + std::shared_ptr default_layout, + client->GetDefaultPjRtLayout(dtype, shape.dims(), + client->addressable_devices().at(0), + MemoryKind())); + xla::Layout transposed_layout = default_layout->xla_layout(); + // We assume that reversing the minor_to_major still gives a valid layout + // for this shape. + absl::c_reverse(*transposed_layout.mutable_minor_to_major()); + layout1 = xla::ifrt::PjRtLayout::Create( + std::make_shared(std::move(transposed_layout))); + } ASSERT_OK_AND_ASSIGN( ArrayRef array1, client->MakeArrayFromHostBuffer( From 1e2aa2dd5f585141a6be7c6a4c9b21e3d962c84c Mon Sep 17 00:00:00 2001 From: Dmitri Latushko Date: Fri, 28 Aug 2026 16:49:29 -0700 Subject: [PATCH 02/12] Clarify in tf.sequence_mask documentation that maxlen must be a compile-time constant when compiling under XLA. PiperOrigin-RevId: 972857042 --- .../kernel_tests/array_ops/array_ops_test.py | 18 ++++++++++++++++++ tensorflow/python/ops/array_ops.py | 4 +++- 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/tensorflow/python/kernel_tests/array_ops/array_ops_test.py b/tensorflow/python/kernel_tests/array_ops/array_ops_test.py index 2fa0dd361ee02f..be29ce381f750e 100644 --- a/tensorflow/python/kernel_tests/array_ops/array_ops_test.py +++ b/tensorflow/python/kernel_tests/array_ops/array_ops_test.py @@ -1650,6 +1650,24 @@ def check_output_dtype(output_dtype): check_output_dtype("float64") check_output_dtype(np.float64) + def testXlaJitCompileWithStaticMaxlen(self): + if not context.executing_eagerly() or not test_util.is_xla_enabled(): + return + + @def_function.function(jit_compile=True) + def fn(lengths): + return array_ops.sequence_mask(lengths, maxlen=5) + + res = fn(constant_op.constant([1, 3, 2])) + self.assertAllEqual( + res, + [ + [True, False, False, False, False], + [True, True, True, False, False], + [True, True, False, False, False], + ], + ) + class ConcatSliceResourceTest(test_util.TensorFlowTestCase): diff --git a/tensorflow/python/ops/array_ops.py b/tensorflow/python/ops/array_ops.py index f84c1f02bd32fd..173a87213ff651 100644 --- a/tensorflow/python/ops/array_ops.py +++ b/tensorflow/python/ops/array_ops.py @@ -4176,7 +4176,9 @@ def sequence_mask(lengths, maxlen=None, dtype=dtypes.bool, name=None): Args: lengths: integer tensor, all its values <= maxlen. maxlen: scalar integer tensor, size of last dimension of returned tensor. - Default is the maximum value in `lengths`. + Default is the maximum value in `lengths`. When compiling with XLA (such + as with `tf.function(jit_compile=True)`), `maxlen` must be explicitly + provided and evaluate to a compile-time constant. dtype: output type of the resulting tensor. name: name of the op. From 96c9a619a0f5d14c0281c839831413899fb2cfe3 Mon Sep 17 00:00:00 2001 From: Maxim Ermilov Date: Fri, 28 Aug 2026 17:15:10 -0700 Subject: [PATCH 03/12] log mlir/triton compile time PiperOrigin-RevId: 972867033 --- .../xla/xla/backends/gpu/codegen/BUILD | 3 + .../codegen/cubin_custom_kernel_compiler.cc | 63 ++++++++++++++----- 2 files changed, 51 insertions(+), 15 deletions(-) diff --git a/third_party/xla/xla/backends/gpu/codegen/BUILD b/third_party/xla/xla/backends/gpu/codegen/BUILD index cfb30514a283cd..94c732c4dc10e1 100644 --- a/third_party/xla/xla/backends/gpu/codegen/BUILD +++ b/third_party/xla/xla/backends/gpu/codegen/BUILD @@ -255,6 +255,7 @@ cc_library( deps = [ ":kernel_compiler", "//xla:future", + "//xla:util", "//xla:xla_proto_cc", "//xla/backends/gpu/codegen/emitters:mlir_kernel_emitter", "//xla/backends/gpu/codegen/kernels:custom_kernel", @@ -274,10 +275,12 @@ cc_library( "@com_google_absl//absl/functional:any_invocable", "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:string_view", "@llvm-project//llvm:Core", "@llvm-project//llvm:OrcJIT", "@llvm-project//llvm:TargetParser", + "@llvm-project//mlir:IR", ], ) diff --git a/third_party/xla/xla/backends/gpu/codegen/cubin_custom_kernel_compiler.cc b/third_party/xla/xla/backends/gpu/codegen/cubin_custom_kernel_compiler.cc index a7dc1a857b3ecd..fa52a302d49cf9 100644 --- a/third_party/xla/xla/backends/gpu/codegen/cubin_custom_kernel_compiler.cc +++ b/third_party/xla/xla/backends/gpu/codegen/cubin_custom_kernel_compiler.cc @@ -23,10 +23,12 @@ limitations under the License. #include "absl/status/status_macros.h" #include "absl/status/statusor.h" +#include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "llvm/ExecutionEngine/Orc/ThreadSafeModule.h" #include "llvm/IR/Module.h" #include "llvm/TargetParser/Triple.h" +#include "mlir/IR/MLIRContext.h" #include "xla/backends/gpu/codegen/emitters/mlir_kernel_emitter.h" #include "xla/backends/gpu/codegen/kernel_compiler.h" #include "xla/backends/gpu/codegen/kernels/custom_kernel.h" @@ -43,11 +45,42 @@ limitations under the License. #include "xla/hlo/ir/hlo_module.h" #include "xla/service/gpu/launch_dimensions.h" #include "xla/stream_executor/device_description.h" +#include "xla/util.h" namespace xla::gpu { - using ::xla::xtile::BlockLevelParameters; +namespace { +absl::StatusOr CompileMlirToLlvmImpl( + const se::DeviceDescription& device, const HloModule& hlo_module, + const std::string& entry_function_name, int unroll_factor, + MlirKernelSource source, mlir::MLIRContext& mlir_context) { + XLA_SCOPED_LOGGING_TIMER_IF( + absl::StrCat("CubinCustomKernelCompiler::CompileMlirToLlvmImpl for ", + hlo_module.name()), + hlo_module.config().debug_options().xla_enable_scoped_logging_timers()); + return gpu::CompileMlirToLlvm(device, hlo_module, entry_function_name, + unroll_factor, mlir_context, std::move(source)); +} + +absl::StatusOr CompileTritonToLlvmImpl( + absl::string_view kernel_name, const HloModule& hlo_module, + const se::DeviceDescription& device_info, + const BlockLevelParameters& block_level_parameters, + const llvm::Triple& target_triple, const std::string& data_layout, + TritonKernelSource triton_source, BorrowedMlirContext borrowed_context, + bool is_xla_fusion) { + XLA_SCOPED_LOGGING_TIMER_IF( + absl::StrCat("CubinCustomKernelCompiler::CompileTritonToLlvmImpl for ", + hlo_module.name()), + hlo_module.config().debug_options().xla_enable_scoped_logging_timers()); + return gpu::CompileTritonToLLVM(kernel_name, hlo_module, device_info, + block_level_parameters, target_triple, + data_layout, std::move(triton_source), + **borrowed_context, is_xla_fusion); +} +} // namespace + xla::Future> CubinCustomKernelCompiler::Compile( Thunk::ThunkInfo thunk_info, LlvmKernelSource kernel_source, const std::string& sanitized_kernel_name, @@ -74,18 +107,18 @@ xla::Future CubinCustomKernelCompiler::CompileMlirToLlvm( const std::string& entry_function_name, int unroll_factor, MlirKernelSource source, BorrowedMlirContext borrowed_context) { if (!thread_pool_) { - return gpu::CompileMlirToLlvm(device, hlo_module, entry_function_name, - unroll_factor, **borrowed_context, - std::move(source)); + return CompileMlirToLlvmImpl(device, hlo_module, entry_function_name, + unroll_factor, std::move(source), + **borrowed_context); } return xla::MakeFutureOn( *thread_pool_->AsExecutor(), [source = std::move(source), device, &hlo_module, entry_function_name, unroll_factor, borrowed_context = std::move(borrowed_context)]() mutable { - return gpu::CompileMlirToLlvm(device, hlo_module, entry_function_name, - unroll_factor, **borrowed_context, - std::move(source)); + return CompileMlirToLlvmImpl(device, hlo_module, entry_function_name, + unroll_factor, std::move(source), + **borrowed_context); }); } @@ -144,10 +177,10 @@ xla::Future CubinCustomKernelCompiler::CompileTritonToLlvm( TritonKernelSource triton_source, BorrowedMlirContext borrowed_context, bool is_xla_fusion) { if (!thread_pool_) { - return gpu::CompileTritonToLLVM(kernel_name, hlo_module, device_info, - block_level_parameters, target_triple, - data_layout, std::move(triton_source), - **borrowed_context, is_xla_fusion); + return CompileTritonToLlvmImpl(kernel_name, hlo_module, device_info, + block_level_parameters, target_triple, + data_layout, std::move(triton_source), + std::move(borrowed_context), is_xla_fusion); } return xla::MakeFutureOn( *thread_pool_->AsExecutor(), @@ -155,10 +188,10 @@ xla::Future CubinCustomKernelCompiler::CompileTritonToLlvm( device_info, block_level_parameters, target_triple, is_xla_fusion, data_layout, borrowed_context = std::move(borrowed_context), triton_source = std::move(triton_source)]() mutable { - return gpu::CompileTritonToLLVM(kernel_name, *hlo_module, device_info, - block_level_parameters, target_triple, - data_layout, std::move(triton_source), - **borrowed_context, is_xla_fusion); + return CompileTritonToLlvmImpl( + kernel_name, *hlo_module, device_info, block_level_parameters, + target_triple, data_layout, std::move(triton_source), + std::move(borrowed_context), is_xla_fusion); }); } From 174f0abd1db70e05b62db6761765c55a02ae8fb6 Mon Sep 17 00:00:00 2001 From: Michael Wong Date: Fri, 28 Aug 2026 17:29:36 -0700 Subject: [PATCH 04/12] Preserve send/recv HloInstruction shape when deserializing from proto. PiperOrigin-RevId: 972872123 --- third_party/xla/xla/hlo/ir/hlo_instruction.cc | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction.cc b/third_party/xla/xla/hlo/ir/hlo_instruction.cc index 840dfc48d04020..be8d15cbc75aee 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instruction.cc +++ b/third_party/xla/xla/hlo/ir/hlo_instruction.cc @@ -546,6 +546,10 @@ absl::StatusOr> HloInstruction::CreateFromProto( case HloOpcode::kSend: instruction = CreateSend(operands(0), operands(1), channel_id, proto.is_host_transfer()); + // CreateSend will create an assumed layout-less u32[] in the output + // shape, so copy over the shape from the proto to ensure no information + // is lost. + *instruction->mutable_shape() = shape; break; case HloOpcode::kSendDone: TF_RET_CHECK(DynCast(operands(0)) != nullptr) @@ -556,6 +560,10 @@ absl::StatusOr> HloInstruction::CreateFromProto( case HloOpcode::kRecv: instruction = CreateRecv(shape.tuple_shapes(0), operands(0), channel_id, proto.is_host_transfer()); + // CreateRecv will create an assumed layout-less u32[] in the output + // shape, so copy over the shape from the proto to ensure no information + // is lost. + *instruction->mutable_shape() = shape; break; case HloOpcode::kRecvDone: TF_RET_CHECK(DynCast(operands(0)) != nullptr) From 2edfe1b45d4345e80eaca49b6b9c5203889e6527 Mon Sep 17 00:00:00 2001 From: Majid Dadashi Date: Fri, 28 Aug 2026 17:30:13 -0700 Subject: [PATCH 05/12] Rename propagate-qsv pass to propagate-qparams and add VHLO quant custom call legalization - Rename PropagateQsvPass to PropagateQParamsPass across passes definitions, headers, implementations, tests, and converter API registration. - Add backward propagation through TransposeOp and ReshapeOp for per-axis quantization parameters. - Add LegalizeVhloQuantCustomCallsPass to legalize vhlo quant custom calls (quant.dequantize, quant.quantize, quant.fake_quant) to stablehlo.custom_call. PiperOrigin-RevId: 972872299 --- tensorflow/compiler/mlir/lite/BUILD | 2 +- tensorflow/compiler/mlir/lite/stablehlo/BUILD | 32 ++- .../legalize-vhlo-quant-custom-calls.mlir | 92 ++++++++ .../legalize_vhlo_quant_custom_calls.cc | 218 ++++++++++++++++++ .../stablehlo/transforms/stablehlo_passes.h | 4 + .../stablehlo/transforms/stablehlo_passes.td | 6 + ...pagate-qsv.mlir => propagate-qparams.mlir} | 2 +- .../compiler/mlir/lite/transforms/passes.h | 6 +- .../compiler/mlir/lite/transforms/passes.td | 10 +- ..._qsv_pass.cc => propagate_qparams_pass.cc} | 121 ++++++++-- 10 files changed, 462 insertions(+), 31 deletions(-) create mode 100644 tensorflow/compiler/mlir/lite/stablehlo/tests/legalize-vhlo-quant-custom-calls.mlir create mode 100644 tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_vhlo_quant_custom_calls.cc rename tensorflow/compiler/mlir/lite/tests/{propagate-qsv.mlir => propagate-qparams.mlir} (99%) rename tensorflow/compiler/mlir/lite/transforms/quantization/{propagate_qsv_pass.cc => propagate_qparams_pass.cc} (73%) diff --git a/tensorflow/compiler/mlir/lite/BUILD b/tensorflow/compiler/mlir/lite/BUILD index 10a286bab5d0b2..f9e50f01ab7481 100644 --- a/tensorflow/compiler/mlir/lite/BUILD +++ b/tensorflow/compiler/mlir/lite/BUILD @@ -1511,7 +1511,7 @@ cc_library( "transforms/prepare_quantize_helper.cc", "transforms/quantization/bias_quantizer_pass.cc", "transforms/quantization/fuse_qdq_pass.cc", - "transforms/quantization/propagate_qsv_pass.cc", + "transforms/quantization/propagate_qparams_pass.cc", "transforms/quantization/quant_utils.cc", "transforms/quantize.cc", "transforms/quantize_variables.cc", diff --git a/tensorflow/compiler/mlir/lite/stablehlo/BUILD b/tensorflow/compiler/mlir/lite/stablehlo/BUILD index e4273397d73a42..1eefbdda3d0e31 100644 --- a/tensorflow/compiler/mlir/lite/stablehlo/BUILD +++ b/tensorflow/compiler/mlir/lite/stablehlo/BUILD @@ -76,6 +76,35 @@ cc_library( alwayslink = 1, ) +cc_library( + name = "legalize_vhlo_quant_custom_calls", + srcs = [ + "transforms/legalize_vhlo_quant_custom_calls.cc", + ], + hdrs = [ + "transforms/stablehlo_passes.h", + "transforms/stablehlo_passes.h.inc", + ], + copts = [ + "-Ithird_party", + ], + deps = [ + ":passes_inc_gen", + "@com_google_absl//absl/strings", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:FuncDialect", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:Support", + "@llvm-project//mlir:TransformUtils", + "@llvm-project//mlir:Transforms", + "@stablehlo//:stablehlo_ops", + "@stablehlo//:vhlo_ops", + "@stablehlo//:vhlo_types", + ], + alwayslink = 1, +) + cc_library( name = "stablehlo_util", srcs = [ @@ -355,7 +384,6 @@ cc_library( "@llvm-project//mlir:Pass", "@llvm-project//mlir:QuantOps", "@llvm-project//mlir:Support", - "@llvm-project//mlir:Transforms", "@xla//xla/mlir_hlo", ], alwayslink = 1, @@ -405,6 +433,7 @@ cc_library( "-Ithird_party", ], deps = [ + ":legalize_stablehlo_custom_call_to_composite", "//tensorflow/compiler/mlir/lite:tensorflow_lite", "@flatbuffers", "@llvm-project//llvm:Support", @@ -937,6 +966,7 @@ tf_cc_binary( ":legalize_stablehlo_custom_call_to_composite", ":legalize_stablehlo_to_vhlo_pass", ":legalize_tf_xla_call_module_to_stablehlo_pass", + ":legalize_vhlo_quant_custom_calls", ":optimize", ":passes_inc_gen", ":prepare_hlo", diff --git a/tensorflow/compiler/mlir/lite/stablehlo/tests/legalize-vhlo-quant-custom-calls.mlir b/tensorflow/compiler/mlir/lite/stablehlo/tests/legalize-vhlo-quant-custom-calls.mlir new file mode 100644 index 00000000000000..2af035360c7b5d --- /dev/null +++ b/tensorflow/compiler/mlir/lite/stablehlo/tests/legalize-vhlo-quant-custom-calls.mlir @@ -0,0 +1,92 @@ +// Copyright 2026 The TensorFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// ============================================================================== +// RUN: odml-to-stablehlo-opt %s -legalize-vhlo-quant-custom-calls -split-input-file | FileCheck %s + +// CHECK-LABEL: func.func @legalize_vhlo_quant_dequantize +func.func @legalize_vhlo_quant_dequantize(%arg0: tensor<4xi8>, %arg1: tensor, %arg2: tensor) -> tensor<4xf32> { + // CHECK-NOT: vhlo.custom_call_v1 + // CHECK: %[[RES:.*]] = stablehlo.custom_call @quant.dequantize(%arg0, %arg1, %arg2) {axis = 0 : i32} : (tensor<4xi8>, tensor, tensor) -> tensor<4xf32> + // CHECK: return %[[RES]] : tensor<4xf32> + %0 = "vhlo.custom_call_v1"(%arg0, %arg1, %arg2) <{ + api_version = #vhlo, + backend_config = #vhlo.string_v1<"">, + call_target_name = #vhlo.string_v1<"quant.dequantize">, + called_computations = #vhlo.array_v1<[]>, + has_side_effect = #vhlo.bool_v1, + operand_layouts = #vhlo.array_v1<[]>, + output_operand_aliases = #vhlo.array_v1<[]>, + result_layouts = #vhlo.array_v1<[]> + }> {axis = #vhlo.integer_v1<0 : i32>} : (tensor<4xi8>, tensor, tensor) -> tensor<4xf32> + return %0 : tensor<4xf32> +} + +// ----- + +// CHECK-LABEL: func.func @legalize_vhlo_quant_quantize +func.func @legalize_vhlo_quant_quantize(%arg0: tensor<4xf32>, %arg1: tensor, %arg2: tensor) -> tensor<4xi8> { + // CHECK-NOT: vhlo.custom_call_v1 + // CHECK: %[[RES:.*]] = stablehlo.custom_call @quant.quantize(%arg0, %arg1, %arg2) : (tensor<4xf32>, tensor, tensor) -> tensor<4xi8> + // CHECK: return %[[RES]] : tensor<4xi8> + %0 = "vhlo.custom_call_v1"(%arg0, %arg1, %arg2) <{ + api_version = #vhlo, + backend_config = #vhlo.string_v1<"">, + call_target_name = #vhlo.string_v1<"quant.quantize">, + called_computations = #vhlo.array_v1<[]>, + has_side_effect = #vhlo.bool_v1, + operand_layouts = #vhlo.array_v1<[]>, + output_operand_aliases = #vhlo.array_v1<[]>, + result_layouts = #vhlo.array_v1<[]> + }> : (tensor<4xf32>, tensor, tensor) -> tensor<4xi8> + return %0 : tensor<4xi8> +} + +// ----- + +// CHECK-LABEL: func.func @legalize_vhlo_quant_fake_quant +func.func @legalize_vhlo_quant_fake_quant(%arg0: tensor<4xf32>, %arg1: tensor, %arg2: tensor) -> tensor<4xf32> { + // CHECK-NOT: vhlo.custom_call_v1 + // CHECK: %[[RES:.*]] = stablehlo.custom_call @quant.fake_quant(%arg0, %arg1, %arg2) {narrow_range = false} : (tensor<4xf32>, tensor, tensor) -> tensor<4xf32> + // CHECK: return %[[RES]] : tensor<4xf32> + %0 = "vhlo.custom_call_v1"(%arg0, %arg1, %arg2) <{ + api_version = #vhlo, + backend_config = #vhlo.string_v1<"">, + call_target_name = #vhlo.string_v1<"quant.fake_quant">, + called_computations = #vhlo.array_v1<[]>, + has_side_effect = #vhlo.bool_v1, + operand_layouts = #vhlo.array_v1<[]>, + output_operand_aliases = #vhlo.array_v1<[]>, + result_layouts = #vhlo.array_v1<[]> + }> {narrow_range = #vhlo.bool_v1} : (tensor<4xf32>, tensor, tensor) -> tensor<4xf32> + return %0 : tensor<4xf32> +} + +// ----- + +// CHECK-LABEL: func.func @keep_other_vhlo_custom_call +func.func @keep_other_vhlo_custom_call(%arg0: tensor<4xf32>) -> tensor<4xf32> { + // CHECK: "vhlo.custom_call_v1"(%arg0) + // CHECK-SAME: call_target_name = #vhlo.string_v1<"other.custom_call"> + %0 = "vhlo.custom_call_v1"(%arg0) <{ + api_version = #vhlo, + backend_config = #vhlo.string_v1<"">, + call_target_name = #vhlo.string_v1<"other.custom_call">, + called_computations = #vhlo.array_v1<[]>, + has_side_effect = #vhlo.bool_v1, + operand_layouts = #vhlo.array_v1<[]>, + output_operand_aliases = #vhlo.array_v1<[]>, + result_layouts = #vhlo.array_v1<[]> + }> : (tensor<4xf32>) -> tensor<4xf32> + return %0 : tensor<4xf32> +} diff --git a/tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_vhlo_quant_custom_calls.cc b/tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_vhlo_quant_custom_calls.cc new file mode 100644 index 00000000000000..9c64b04e36f27a --- /dev/null +++ b/tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_vhlo_quant_custom_calls.cc @@ -0,0 +1,218 @@ +/* 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 "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project +#include "mlir/IR/Attributes.h" // from @llvm-project +#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project +#include "mlir/IR/BuiltinOps.h" // from @llvm-project +#include "mlir/IR/BuiltinTypes.h" // from @llvm-project +#include "mlir/IR/MLIRContext.h" // from @llvm-project +#include "mlir/IR/Operation.h" // from @llvm-project +#include "mlir/IR/PatternMatch.h" // from @llvm-project +#include "mlir/IR/Types.h" // from @llvm-project +#include "mlir/IR/Value.h" // from @llvm-project +#include "mlir/Pass/Pass.h" // from @llvm-project +#include "mlir/Pass/PassRegistry.h" // from @llvm-project +#include "mlir/Support/LLVM.h" // from @llvm-project +#include "mlir/Support/LogicalResult.h" // from @llvm-project +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project +#include "stablehlo/dialect/StablehloOps.h" // from @stablehlo // build_cleaner: keep +#include "stablehlo/dialect/VhloOps.h" // from @stablehlo // build_cleaner: keep +#include "stablehlo/dialect/VhloTypes.h" // from @stablehlo +#include "tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.h" + +namespace mlir::odml { + +#define GEN_PASS_DEF_LEGALIZEVHLOQUANTCUSTOMCALLSPASS +#include "tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.h.inc" + +namespace { + +class VhloToStablehloTypeConverter : public vhlo::VhloTypeConverter { + public: + VhloToStablehloTypeConverter() + : vhlo::VhloTypeConverter(/*allowOtherDialects=*/true) { + addConversion([](Type type) -> Type { return type; }); + addConversion([](vhlo::TokenV1Type token) -> Type { + return stablehlo::TokenType::get(token.getContext()); + }); + addVhloToBuiltinConversions(); + } + + Attribute convertEncoding(Attribute attr) const final { + if (auto vhloAttr = + mlir::dyn_cast_or_null(attr)) { + return stablehlo::TypeExtensionsAttr::get(vhloAttr.getContext(), + vhloAttr.getBounds()); + } + return attr; + } +}; + +Attribute ConvertVhloAttrToBuiltin(Attribute attr) { + if (!attr) return {}; + if (auto vhlo_str = mlir::dyn_cast(attr)) { + return StringAttr::get(attr.getContext(), vhlo_str.getValue()); + } + if (auto vhlo_bool = mlir::dyn_cast(attr)) { + return BoolAttr::get(attr.getContext(), vhlo_bool.getValue()); + } + if (auto vhlo_int = mlir::dyn_cast(attr)) { + VhloToStablehloTypeConverter type_converter; + Type type = type_converter.convertType(vhlo_int.getType()); + if (!type) type = vhlo_int.getType(); + return IntegerAttr::get(type, vhlo_int.getValue()); + } + if (auto vhlo_float = mlir::dyn_cast(attr)) { + VhloToStablehloTypeConverter type_converter; + Type type = type_converter.convertType(vhlo_float.getType()); + if (!type) type = vhlo_float.getType(); + return FloatAttr::get(type, vhlo_float.getValue()); + } + if (auto vhlo_tensor = mlir::dyn_cast(attr)) { + return vhlo_tensor.getData(); + } + return attr; +} + +struct LegalizeVhloQuantCustomCallPattern : public RewritePattern { + explicit LegalizeVhloQuantCustomCallPattern(MLIRContext* context) + : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {} + + LogicalResult matchAndRewrite(Operation* op, + PatternRewriter& rewriter) const override { + StringRef op_name = op->getName().getStringRef(); + if (op_name != "vhlo.custom_call_v1" && op_name != "vhlo.custom_call") { + return failure(); + } + + Attribute raw_target = op->getAttr("call_target_name"); + if (!raw_target) return failure(); + + StringRef target_name; + if (auto attr = mlir::dyn_cast(raw_target)) { + target_name = attr.getValue(); + } else if (auto attr = mlir::dyn_cast(raw_target)) { + target_name = attr.getValue(); + } + + if (target_name != "quant.dequantize" && target_name != "quant.quantize" && + target_name != "quant.fake_quant") { + return failure(); + } + + VhloToStablehloTypeConverter type_converter; + + SmallVector result_types; + for (Type t : op->getResultTypes()) { + Type conv = type_converter.convertType(t); + result_types.push_back(conv ? conv : t); + } + + SmallVector operands; + for (Value val : op->getOperands()) { + Type conv = type_converter.convertType(val.getType()); + if (conv && conv != val.getType()) { + val = + rewriter.create(op->getLoc(), conv, val) + .getResult(0); + } + operands.push_back(val); + } + + SmallVector new_attrs; + new_attrs.push_back(rewriter.getNamedAttr( + "call_target_name", rewriter.getStringAttr(target_name))); + + static const char* const kIntrinsicAttrs[] = {"api_version", + "backend_config", + "call_target_name", + "called_computations", + "has_side_effect", + "operand_layouts", + "output_operand_aliases", + "result_layouts", + "result_tilings"}; + + for (NamedAttribute attr : op->getAttrs()) { + StringRef name = attr.getName().strref(); + bool is_intrinsic = false; + for (const char* kAttr : kIntrinsicAttrs) { + if (name == kAttr) { + is_intrinsic = true; + break; + } + } + if (is_intrinsic) continue; + + Attribute builtin_val = ConvertVhloAttrToBuiltin(attr.getValue()); + if (builtin_val) { + new_attrs.push_back(rewriter.getNamedAttr(name, builtin_val)); + } + } + + auto new_op = rewriter.create( + op->getLoc(), result_types, operands, new_attrs); + + if (new_op->getNumResults() != op->getNumResults()) { + return failure(); + } + + SmallVector replacement_vals; + for (unsigned i = 0; i < op->getNumResults(); ++i) { + Value new_res = new_op.getResult(i); + Type orig_type = op->getResult(i).getType(); + if (new_res.getType() != orig_type) { + new_res = rewriter + .create(op->getLoc(), + orig_type, new_res) + .getResult(0); + } + replacement_vals.push_back(new_res); + } + + rewriter.replaceOp(op, replacement_vals); + return success(); + } +}; + +class LegalizeVhloQuantCustomCallsPass + : public impl::LegalizeVhloQuantCustomCallsPassBase< + LegalizeVhloQuantCustomCallsPass> { + public: + void runOnOperation() override { + ModuleOp module = getOperation(); + MLIRContext* context = &getContext(); + RewritePatternSet patterns(context); + patterns.add(context); + if (failed(applyPatternsGreedily(module, std::move(patterns)))) { + signalPassFailure(); + } + } +}; + +} // namespace + +std::unique_ptr> +CreateLegalizeVhloQuantCustomCallsPass() { + return std::make_unique(); +} + +static PassRegistration pass; + +} // namespace mlir::odml diff --git a/tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.h b/tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.h index 3328db76764ce7..7ec56df1e24504 100644 --- a/tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.h +++ b/tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.h @@ -67,6 +67,10 @@ void PopulateLegalizeHloToTfPatterns(RewritePatternSet* patterns, // Drops vhlo/stablehlo custom calls targeting 'shape_assertion'. std::unique_ptr> CreateDropShapeAssertionsPass(); +// Legalizes vhlo custom calls for quantization ops to stablehlo.custom_call. +std::unique_ptr> +CreateLegalizeVhloQuantCustomCallsPass(); + #define GEN_PASS_DECL #include "tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.h.inc" diff --git a/tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.td b/tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.td index ec2145bf990f60..8618ce4784b474 100644 --- a/tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.td +++ b/tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.td @@ -192,3 +192,9 @@ def DropShapeAssertionsPass : Pass<"drop-shape-assertions", "ModuleOp"> { let constructor = "mlir::odml::CreateDropShapeAssertionsPass()"; let dependentDialects = ["mlir::vhlo::VhloDialect", "mlir::stablehlo::StablehloDialect"]; } + +def LegalizeVhloQuantCustomCallsPass : Pass<"legalize-vhlo-quant-custom-calls", "ModuleOp"> { + let summary = "Legalizes vhlo custom calls for quantization ops to stablehlo.custom_call."; + let constructor = "mlir::odml::CreateLegalizeVhloQuantCustomCallsPass()"; + let dependentDialects = ["mlir::vhlo::VhloDialect", "mlir::stablehlo::StablehloDialect"]; +} diff --git a/tensorflow/compiler/mlir/lite/tests/propagate-qsv.mlir b/tensorflow/compiler/mlir/lite/tests/propagate-qparams.mlir similarity index 99% rename from tensorflow/compiler/mlir/lite/tests/propagate-qsv.mlir rename to tensorflow/compiler/mlir/lite/tests/propagate-qparams.mlir index 2ffdbd42524adc..7fa7a19c720003 100644 --- a/tensorflow/compiler/mlir/lite/tests/propagate-qsv.mlir +++ b/tensorflow/compiler/mlir/lite/tests/propagate-qparams.mlir @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================== -// RUN: litert-opt %s -tfl-propagate-qsv | FileCheck %s +// RUN: litert-opt %s -tfl-propagate-qparams | FileCheck %s // CHECK-LABEL: concat diff --git a/tensorflow/compiler/mlir/lite/transforms/passes.h b/tensorflow/compiler/mlir/lite/transforms/passes.h index c969f66f7c7191..84b7476d569fed 100644 --- a/tensorflow/compiler/mlir/lite/transforms/passes.h +++ b/tensorflow/compiler/mlir/lite/transforms/passes.h @@ -124,9 +124,9 @@ std::unique_ptr> CreateDefaultQuantizePass(); std::unique_ptr> CreateLowerQuantAnnotationsPass(); -// Creates an instance of the TFLite PropagateQsv pass which propagates scale -// and zero point (QSV) information through the graph. -std::unique_ptr> CreatePropagateQsvPass(); +// Creates an instance of the TFLite PropagateQParams pass which propagates +// scale and zero point (quantization parameters) through the graph. +std::unique_ptr> CreatePropagateQParamsPass(); std::unique_ptr> CreateBiasQuantizerPass(); diff --git a/tensorflow/compiler/mlir/lite/transforms/passes.td b/tensorflow/compiler/mlir/lite/transforms/passes.td index 2592531b83c031..5aab64ac05b4d9 100644 --- a/tensorflow/compiler/mlir/lite/transforms/passes.td +++ b/tensorflow/compiler/mlir/lite/transforms/passes.td @@ -354,13 +354,13 @@ def LowerQuantAnnotationsPass : Pass<"tfl-lower-quant-annotations", "mlir::Modul ]; } -def PropagateQsvPass : Pass<"tfl-propagate-qsv", "mlir::ModuleOp"> { - let summary = "Propagates Quantization Scale/Value (QSV) information through the graph."; +def PropagateQParamsPass : Pass<"tfl-propagate-qparams", "mlir::ModuleOp"> { + let summary = "Propagates Quantization Parameters (scale and zero point) information through the graph."; let description = [{ - This transformation pass propagates the QSV data across operations in the - TensorFlow Lite dialect. + This transformation pass propagates the quantization parameters across + operations in the TensorFlow Lite dialect. }]; - let constructor = "CreatePropagateQsvPass()"; + let constructor = "CreatePropagateQParamsPass()"; let dependentDialects = [ "TFL::TensorFlowLiteDialect", "mlir::quant::QuantDialect" diff --git a/tensorflow/compiler/mlir/lite/transforms/quantization/propagate_qsv_pass.cc b/tensorflow/compiler/mlir/lite/transforms/quantization/propagate_qparams_pass.cc similarity index 73% rename from tensorflow/compiler/mlir/lite/transforms/quantization/propagate_qsv_pass.cc rename to tensorflow/compiler/mlir/lite/transforms/quantization/propagate_qparams_pass.cc index e4820aa0997f38..99f925ca159e93 100644 --- a/tensorflow/compiler/mlir/lite/transforms/quantization/propagate_qsv_pass.cc +++ b/tensorflow/compiler/mlir/lite/transforms/quantization/propagate_qparams_pass.cc @@ -13,7 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -// This transformation pass propagates QSV information through the model. +// This transformation pass propagates quantization parameters through the +// model. #include #include @@ -24,7 +25,6 @@ limitations under the License. #include "absl/container/flat_hash_set.h" #include "absl/status/statusor.h" -#include "llvm/Support/ErrorHandling.h" #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project #include "mlir/Dialect/Quant/IR/Quant.h" // from @llvm-project #include "mlir/Dialect/Quant/IR/QuantTypes.h" // from @llvm-project @@ -57,15 +57,15 @@ namespace mlir { namespace TFL { namespace { -#define GEN_PASS_DEF_PROPAGATEQSVPASS +#define GEN_PASS_DEF_PROPAGATEQPARAMSPASS #include "tensorflow/compiler/mlir/lite/transforms/passes.h.inc" //-------------------------------------------------------------------------===// // Helper Functions //===----------------------------------------------------------------------===// -// Returns true if any of the op's operands are per-axis quantized. -bool HasPerAxisQuantizedOperand(mlir::Operation* op) { +// Returns true if any of the op's operands or results are per-axis quantized. +bool HasPerAxisQuantizedValue(mlir::Operation* op) { for (const auto& operand : op->getOperands()) { auto qtype = GetQTypeFromDefiningDequantize(operand); if (qtype.has_value() && @@ -73,13 +73,20 @@ bool HasPerAxisQuantizedOperand(mlir::Operation* op) { return true; } } + for (const auto& result : op->getResults()) { + auto qtype = GetQTypeFromConsumingQuantize(result); + if (qtype.has_value() && + dyn_cast(*qtype)) { + return true; + } + } return false; } // Propagates the quantized type `qtype` to all float operands and results of // `same_scales_op` by inserting QDQ pairs. This is only used for ops that have // the SameScalesOpInterface. -LogicalResult PropagateQsvAcrossOperandsAndResults( +LogicalResult PropagateQParamsAcrossOperandsAndResults( SameScalesOpInterface same_scales_op, quant::QuantizedType qtype, PatternRewriter& rewriter) { mlir::Operation* op = same_scales_op.getOperation(); @@ -164,8 +171,8 @@ LogicalResult GetQuantDimensionAfterTranspose(TFL::TransposeOp transpose_op, // Find what the quantized dimension has been transposed to const auto it = std::find(axes.begin(), axes.end(), quant_dim); if (it == axes.end()) { - llvm_unreachable( - "quantized dimension should be present in a valid permutation"); + return rewriter.notifyMatchFailure(transpose_op, + "quantized dimension not found in perm"); } new_quant_dim = std::distance(axes.begin(), it); return success(); @@ -185,6 +192,33 @@ class PropagateReshapedPerAxisQuantDim std::optional qtype = GetQTypeFromDefiningDequantize(reshape_op.getOperand(0)); if (!qtype.has_value()) { + // Backward propagation: if result is quantized and input is not + std::optional out_qtype = + GetQTypeFromConsumingQuantize(reshape_op.getResult()); + if (out_qtype.has_value()) { + if (auto per_axis_quant = + dyn_cast(*out_qtype)) { + absl::StatusOr in_quant_dim = GetQuantDimensionAfterReshape( + reshape_op.getType().getShape(), + reshape_op.getInput().getType().getShape(), + per_axis_quant.getQuantizedDimension()); + if (in_quant_dim.ok()) { + auto new_element_type = + mlir::quant::UniformQuantizedPerAxisType::getChecked( + reshape_op.getLoc(), per_axis_quant.getFlags(), + per_axis_quant.getStorageType(), + per_axis_quant.getExpressedType(), + per_axis_quant.getScales(), per_axis_quant.getZeroPoints(), + *in_quant_dim, per_axis_quant.getStorageTypeMin(), + per_axis_quant.getStorageTypeMax()); + if (failed(InsertQDQ(reshape_op.getOperand(0), new_element_type, + rewriter, reshape_op))) { + return failure(); + } + return success(); + } + } + } return rewriter.notifyMatchFailure(reshape_op, "input is not a dequantize op"); } @@ -238,6 +272,50 @@ class PropagateTransposedPerAxisQuantDim std::optional qtype = GetQTypeFromDefiningDequantize(transpose_op.getOperand(0)); if (!qtype.has_value()) { + // Backward propagation: if result is quantized and input is not + std::optional out_qtype = + GetQTypeFromConsumingQuantize(transpose_op.getResult()); + if (out_qtype.has_value()) { + if (auto per_axis_quant = + dyn_cast(*out_qtype)) { + DenseIntElementsAttr perm; + if (matchPattern(transpose_op.getPerm(), m_Constant(&perm))) { + auto input_type = + mlir::cast(transpose_op.getInput().getType()); + SmallVector axes; + axes.reserve(perm.getNumElements()); + for (const auto& axis_int : perm.getValues()) { + int64_t axis = axis_int.getSExtValue(); + if (axis < 0) { + axis += input_type.getRank(); + } + if (axis < 0 || + (input_type.hasRank() && axis >= input_type.getRank())) { + continue; + } + axes.push_back(axis); + } + int out_quant_dim = per_axis_quant.getQuantizedDimension(); + if (out_quant_dim >= 0 && out_quant_dim < axes.size()) { + int in_quant_dim = axes[out_quant_dim]; + auto new_element_type = + mlir::quant::UniformQuantizedPerAxisType::getChecked( + transpose_op.getLoc(), per_axis_quant.getFlags(), + per_axis_quant.getStorageType(), + per_axis_quant.getExpressedType(), + per_axis_quant.getScales(), + per_axis_quant.getZeroPoints(), in_quant_dim, + per_axis_quant.getStorageTypeMin(), + per_axis_quant.getStorageTypeMax()); + if (failed(InsertQDQ(transpose_op.getOperand(0), new_element_type, + rewriter, transpose_op))) { + return failure(); + } + return success(); + } + } + } + } return rewriter.notifyMatchFailure(transpose_op, "input is not a dequantize op"); } @@ -278,7 +356,8 @@ class PropagateTransposedPerAxisQuantDim } }; -class PropagateQsv : public OpInterfaceRewritePattern { +class PropagateQParams + : public OpInterfaceRewritePattern { using OpInterfaceRewritePattern< SameScalesOpInterface>::OpInterfaceRewritePattern; @@ -287,7 +366,7 @@ class PropagateQsv : public OpInterfaceRewritePattern { // The per-axis quantized ops that don't directly transfer the quantized // types from input to output (e.g. TransposeOp, ReshapeOp), need dedicated // propagation patterns. - if (!op.RequiredSameQuantizedAxes() && HasPerAxisQuantizedOperand(op)) { + if (!op.RequiredSameQuantizedAxes() && HasPerAxisQuantizedValue(op)) { return rewriter.notifyMatchFailure( op, "requires dedicated propagation pattern."); } @@ -296,7 +375,8 @@ class PropagateQsv : public OpInterfaceRewritePattern { if (!propagated_type) { return rewriter.notifyMatchFailure(op, "No propagated type found."); } - return PropagateQsvAcrossOperandsAndResults(op, *propagated_type, rewriter); + return PropagateQParamsAcrossOperandsAndResults(op, *propagated_type, + rewriter); } }; @@ -304,9 +384,10 @@ class PropagateQsv : public OpInterfaceRewritePattern { // Pass Definition //===----------------------------------------------------------------------===// -struct PropagateQsvPass : public impl::PropagateQsvPassBase { +struct PropagateQParamsPass + : public impl::PropagateQParamsPassBase { public: - MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PropagateQsvPass) + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PropagateQParamsPass) void runOnOperation() override; }; @@ -316,12 +397,12 @@ struct PropagateQsvPass : public impl::PropagateQsvPassBase { //===----------------------------------------------------------------------===// #include "tensorflow/compiler/mlir/lite/transforms/quantization/generated_strict_quantize.inc" -void PropagateQsvPass::runOnOperation() { +void PropagateQParamsPass::runOnOperation() { MLIRContext* ctx = &getContext(); mlir::ModuleOp module = getOperation(); RewritePatternSet patterns(ctx); - patterns.add(ctx); + patterns.add(ctx); // Dedicated propagation patterns. patterns.add> CreatePropagateQsvPass() { - return std::make_unique(); +std::unique_ptr> CreatePropagateQParamsPass() { + return std::make_unique(); } } // namespace TFL From a4a8ad3d388693d94e6824b9bffc4edb64b20436 Mon Sep 17 00:00:00 2001 From: Parker Schuh Date: Fri, 28 Aug 2026 18:06:41 -0700 Subject: [PATCH 06/12] Remove references to StreamExecutorGpuClient. PiperOrigin-RevId: 972887119 --- .../core/common_runtime/gpu/gpu_device.cc | 20 +++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/tensorflow/core/common_runtime/gpu/gpu_device.cc b/tensorflow/core/common_runtime/gpu/gpu_device.cc index 255dcd6297fc46..b58965dc1f5361 100644 --- a/tensorflow/core/common_runtime/gpu/gpu_device.cc +++ b/tensorflow/core/common_runtime/gpu/gpu_device.cc @@ -1739,22 +1739,20 @@ Status BaseGPUDeviceFactory::CreateDevices( : std::make_optional(allowed_devices))); bool should_create_new_pjrt_client = true; - xla::StreamExecutorGpuClient* pjrt_se_client = nullptr; auto obtained_pjrt_client = GetPjRtClient(DeviceType(DEVICE_GPU)); if (obtained_pjrt_client.ok()) { - pjrt_se_client = - absl::down_cast(*obtained_pjrt_client); // TODO(b/291943099): This check may not be enough because the virtual // device options can change while the device count remains the same. // However, it's most likely that in real use cases, CreateDevices() won't // be called more than once with different options being set. If such use // cases exist we may need to update the check here. - if (pjrt_se_client->addressable_device_count() == tf_device_specs.size()) { + if ((*obtained_pjrt_client)->addressable_device_count() == + tf_device_specs.size()) { should_create_new_pjrt_client = false; } else { LOG(WARNING) << "A PjRt GPU Client was previously created, but the " "addressable device count: " - << pjrt_se_client->addressable_device_count() + << (*obtained_pjrt_client)->addressable_device_count() << " is not equal to tf_device_specs size: " << tf_device_specs.size() << ". This usually only happens in unit tests and we will " @@ -1844,9 +1842,15 @@ Status BaseGPUDeviceFactory::CreateDevices( VLOG(3) << "should_create_new_pjrt_client=" << should_create_new_pjrt_client << " for device ordinal " << di << ". Re-using local_device_state"; - auto* pjrt_se_client = - absl::down_cast(*obtained_pjrt_client); - local_device_state = &(pjrt_se_client->device_state(di)); + auto* pjrt_se_client = absl::down_cast( + absl::down_cast(*obtained_pjrt_client) + ->raw_client()); + local_device_state = pjrt_se_client->device_state(xla::LocalDeviceId(di)); + if (!local_device_state) { + return absl::InternalError(absl::StrCat( + "GPU local device state for tf_device_id: ", tf_device_id.value(), + " does not exist.")); + } } // CreateGPUDevice sets stream to `gpu_allocator` and preallocates From bfd5c52f450f54263ee4d3f543f2b2c6572f83cd Mon Sep 17 00:00:00 2001 From: Bhatu Date: Fri, 28 Aug 2026 18:40:04 -0700 Subject: [PATCH 07/12] Clamp propagated constraints across convert operations to the operand's type range to prevent Inf/NaN. During backward constraint propagation, `kConvert` and `kBitcastConvert` operations were passing output constraints through identically to `kCopy` without checking whether the operand has a narrower domain. When downstream 16-bit operations constrained the output to `[-65504, 65504]`, an `f8e4m3fn` operand received this out-of-range interval. Because `f8e4m3fn` has a maximum finite range of `[-448, 448]` and lacks an infinity representation, values outside this range saturated to NaN during random literal generation. This change separates `kConvert` and `kBitcastConvert` in `ConstraintPropagator` to intersect incoming constraints with the finite domain of the operand type, ensuring input parameters are always constrained within their type's representable bounds. PiperOrigin-RevId: 972897290 --- .../xla/xla/tests/constraint_propagator.cc | 42 ++++++++++++++++++- .../xla/tests/constraint_propagator_test.cc | 42 +++++++++++++++++++ 2 files changed, 82 insertions(+), 2 deletions(-) diff --git a/third_party/xla/xla/tests/constraint_propagator.cc b/third_party/xla/xla/tests/constraint_propagator.cc index 3fd5b3b2426b8e..e47a042be3cc70 100644 --- a/third_party/xla/xla/tests/constraint_propagator.cc +++ b/third_party/xla/xla/tests/constraint_propagator.cc @@ -34,6 +34,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/literal.h" +#include "xla/primitive_util.h" #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/tests/constraint_state.h" @@ -284,6 +285,31 @@ void Seed16BitFloatingInstruction( } } +// Returns the finite representable range [lowest, max] for a given +// PrimitiveType as a ConstraintInterval. +std::optional GetTypeFiniteDomain(PrimitiveType type) { + if (type == BF16 || type == F16) { + // 65504.0 is the maximum finite value of FP16. We deliberately bound BF16 + // to 65504.0 as in Seed16BitFloatingInstruction to prevent precision loss. + return ConstraintInterval{-65504.0, 65504.0, false}; + } + return primitive_util::PrimitiveTypeSwitch>( + [&](auto primitive_type_constant) -> std::optional { + if constexpr (primitive_util::IsFloatingPointType( + primitive_type_constant) || + primitive_util::IsIntegralType(primitive_type_constant) || + primitive_type_constant == PRED) { + using NativeT = primitive_util::NativeTypeOf; + return ConstraintInterval{ + static_cast(std::numeric_limits::lowest()), + static_cast(std::numeric_limits::max()), + /*exclude_zero=*/false}; + } + return std::nullopt; + }, + type); +} + // Finds the maximum magnitude M such that the symmetric interval [-M, M] // is contained inside output interval [-L, R] (where L, R > 0). // @@ -862,8 +888,6 @@ absl::Status ConstraintPropagator::PropagateConstraintsExact( break; } case HloOpcode::kBitcast: - case HloOpcode::kBitcastConvert: - case HloOpcode::kConvert: case HloOpcode::kCopy: case HloOpcode::kDynamicReshape: case HloOpcode::kReducePrecision: @@ -871,6 +895,20 @@ absl::Status ConstraintPropagator::PropagateConstraintsExact( case HloOpcode::kTranspose: states_[instruction->operand(0)].Merge(output_state); break; + case HloOpcode::kBitcastConvert: + case HloOpcode::kConvert: { + PrimitiveType operand_type = + instruction->operand(0)->shape().element_type(); + ConstraintState operand_state = output_state; + if (!output_state.GetConstraintInterval().IsUnconstrained()) { + if (std::optional type_bound = + GetTypeFiniteDomain(operand_type)) { + operand_state.AddConstraint(*type_bound); + } + } + states_[instruction->operand(0)].Merge(operand_state); + break; + } case HloOpcode::kReverse: { states_[instruction->operand(0)].AddConstraint(output_interval); StructuralConstraints sc = output_structural; diff --git a/third_party/xla/xla/tests/constraint_propagator_test.cc b/third_party/xla/xla/tests/constraint_propagator_test.cc index 3db05b09e4aed0..0b937c57f66212 100644 --- a/third_party/xla/xla/tests/constraint_propagator_test.cc +++ b/third_party/xla/xla/tests/constraint_propagator_test.cc @@ -1269,5 +1269,47 @@ ENTRY %main { EXPECT_EQ(GetMaxAddReductionElementsForExp(propagator, exp_reduced), 256); EXPECT_EQ(GetMaxAddReductionElementsForExp(propagator, exp_unreduced), 1); } + +TEST_F(ConstraintPropagatorTest, ConvertClampsIntervalToOperandTypeDomain) { + const char* hlo = R"( +HloModule TestModule +ENTRY main { + x = f8e4m3fn[] parameter(0) + ROOT root = bf16[] convert(x) +} +)"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo)); + ASSERT_OK_AND_ASSIGN(auto states, ConstraintPropagator::Run(*module)); + + auto x_int = states[module->entry_computation()->parameter_instruction(0)] + .GetConstraintInterval(); + + // Root bf16 is seeded with [-65504.0, 65504.0]. + // Backward propagation through convert must clamp to f8e4m3fn finite domain + // [-448.0, 448.0]. + EXPECT_DOUBLE_EQ(x_int.min, -448.0); + EXPECT_DOUBLE_EQ(x_int.max, 448.0); +} + +TEST_F(ConstraintPropagatorTest, ConvertClampsIntervalToIntegerTypeDomain) { + const char* hlo = R"( +HloModule TestModule +ENTRY main { + x = s8[] parameter(0) + ROOT root = bf16[] convert(x) +} +)"; + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo)); + ASSERT_OK_AND_ASSIGN(auto states, ConstraintPropagator::Run(*module)); + + auto x_int = states[module->entry_computation()->parameter_instruction(0)] + .GetConstraintInterval(); + + // Root bf16 is seeded with [-65504.0, 65504.0]. + // Backward propagation through convert must clamp to s8 domain [-128.0, + // 127.0]. + EXPECT_DOUBLE_EQ(x_int.min, -128.0); + EXPECT_DOUBLE_EQ(x_int.max, 127.0); +} } // namespace } // namespace xla From 6eefa0c24afc11041f07f30c1ace6280b13df994 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Fri, 28 Aug 2026 19:15:44 -0700 Subject: [PATCH 08/12] Automated Code Change PiperOrigin-RevId: 972907865 --- .../xla/xla/backends/cpu/nanort/ifrt_client.cc | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/third_party/xla/xla/backends/cpu/nanort/ifrt_client.cc b/third_party/xla/xla/backends/cpu/nanort/ifrt_client.cc index d8dd838b317348..71168862337e01 100644 --- a/third_party/xla/xla/backends/cpu/nanort/ifrt_client.cc +++ b/third_party/xla/xla/backends/cpu/nanort/ifrt_client.cc @@ -574,7 +574,7 @@ class NanoArray final : public NanoValue { OwnedDataPtr owned_data_; }; -ABSL_ATTRIBUTE_UNUSED char NanoArray::ID = 'A'; // NOLINT +[[maybe_unused]] char NanoArray::ID = 'A'; // NOLINT // Sharded array implementation. Represents an array that should be assembled // from multiple arrays, but we aren't sure how to assemble it yet. @@ -693,7 +693,7 @@ class ShardedNanoArray final : public NanoValue { return Ready(Internal("Cannot copy sharded array to host buffer.")); } - ABSL_ATTRIBUTE_UNUSED static char ID; // NOLINT + [[maybe_unused]] static char ID; // NOLINT private: ShardedNanoArray(NanoIfrtClient* client, ifrt::DType dtype, @@ -838,7 +838,7 @@ class NanoTuple final : public NanoValue { std::vector values_; }; -ABSL_ATTRIBUTE_UNUSED char NanoTuple::ID = 'T'; // NOLINT +[[maybe_unused]] char NanoTuple::ID = 'T'; // NOLINT // Executable implementation. class NanoExecutable final @@ -1240,7 +1240,7 @@ class NanoExecutable final const xla::ifrt::UserContextRef user_context_; }; -ABSL_ATTRIBUTE_UNUSED char NanoExecutable::ID = 'E'; // NOLINT +[[maybe_unused]] char NanoExecutable::ID = 'E'; // NOLINT // Compiler implementation. class NanoCompiler final @@ -1280,7 +1280,7 @@ class NanoCompiler final NanoIfrtClient* client_; }; -ABSL_ATTRIBUTE_UNUSED char NanoCompiler::ID = 'C'; // NOLINT +[[maybe_unused]] char NanoCompiler::ID = 'C'; // NOLINT // Memory implementation. There is only one address space so this doesn't do // much. @@ -1310,7 +1310,7 @@ class NanoMemory final NanoIfrtClient* client_; }; -ABSL_ATTRIBUTE_UNUSED char NanoMemory::ID = 'M'; // NOLINT +[[maybe_unused]] char NanoMemory::ID = 'M'; // NOLINT // Device implementation. There is only one device so this doesn't do much. class NanoDevice final @@ -1356,7 +1356,7 @@ class NanoDevice final ifrt::Memory* memory_; }; -ABSL_ATTRIBUTE_UNUSED char NanoDevice::ID = 'D'; // NOLINT +[[maybe_unused]] char NanoDevice::ID = 'D'; // NOLINT } // namespace From eb11ee575ed4b109f6fc83c0b550ce590d6e57d9 Mon Sep 17 00:00:00 2001 From: Ionel Gog Date: Fri, 28 Aug 2026 19:45:10 -0700 Subject: [PATCH 09/12] [IFRT IR] Include the MLIR location in the user contexts created for IRRT IR ops PiperOrigin-RevId: 972917231 --- .../xla/python/ifrt/ir/program_interpreter.cc | 20 ++++++++----------- 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/third_party/xla/xla/python/ifrt/ir/program_interpreter.cc b/third_party/xla/xla/python/ifrt/ir/program_interpreter.cc index 944cf95a33d54a..2bff7853184ec3 100644 --- a/third_party/xla/xla/python/ifrt/ir/program_interpreter.cc +++ b/third_party/xla/xla/python/ifrt/ir/program_interpreter.cc @@ -362,9 +362,8 @@ struct CallLoadedExecutableOpState { VLOG(3) << pretty_print; ifrt::UserContextRef new_context = - env.set_op_user_contexts - ? ifrt::BasicUserContext::Create("Execute program op") - : ifrt::UserContextScope::current(); + env.set_op_user_contexts ? ifrt::BasicUserContext::Create(pretty_print) + : ifrt::UserContextScope::current(); ifrt::UserContextScope context_scope(std::move(new_context)); ExecuteOptions options = execute_options; @@ -571,9 +570,8 @@ struct RemapArraysOpState { VLOG(3) << pretty_print; ifrt::UserContextRef new_context = - env.set_op_user_contexts - ? ifrt::BasicUserContext::Create("RemapArrays program op") - : ifrt::UserContextScope::current(); + env.set_op_user_contexts ? ifrt::BasicUserContext::Create(pretty_print) + : ifrt::UserContextScope::current(); ifrt::UserContextScope context_scope(std::move(new_context)); std::vector inputs; @@ -744,9 +742,8 @@ struct BitcastArraysOpState { VLOG(3) << pretty_print; ifrt::UserContextRef new_context = - env.set_op_user_contexts - ? ifrt::BasicUserContext::Create("BitcastArrays program op") - : ifrt::UserContextScope::current(); + env.set_op_user_contexts ? ifrt::BasicUserContext::Create(pretty_print) + : ifrt::UserContextScope::current(); ifrt::UserContextScope context_scope(std::move(new_context)); std::vector inputs; @@ -876,9 +873,8 @@ struct CopyArraysOpState { VLOG(3) << pretty_print; ifrt::UserContextRef new_context = - env.set_op_user_contexts - ? ifrt::BasicUserContext::Create("CopyArrays program op") - : ifrt::UserContextScope::current(); + env.set_op_user_contexts ? ifrt::BasicUserContext::Create(pretty_print) + : ifrt::UserContextScope::current(); ifrt::UserContextScope context_scope(std::move(new_context)); std::vector inputs; From 2ba1ec5f1687f285108dbbe2b29d1cb6496f2a7e Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Fri, 28 Aug 2026 22:53:34 -0700 Subject: [PATCH 10/12] Automated Code Change PiperOrigin-RevId: 972973724 --- tensorflow/dtensor/mlir/expansions/argmax_spmd_expander.cc | 2 ++ .../dtensor/mlir/expansions/broadcast_to_spmd_expander.cc | 1 + .../dtensor/mlir/expansions/control_flow_spmd_expander.cc | 1 + tensorflow/dtensor/mlir/expansions/conv_spmd_expander.cc | 1 + tensorflow/dtensor/mlir/expansions/cumsum_spmd_expander.cc | 2 +- .../dtensor/mlir/expansions/dataparallel_spmd_expander.cc | 1 + tensorflow/dtensor/mlir/expansions/einsum_spmd_expander.cc | 1 + tensorflow/dtensor/mlir/expansions/elementwise_spmd_expander.cc | 1 + tensorflow/dtensor/mlir/expansions/fill_spmd_expander.cc | 2 ++ tensorflow/dtensor/mlir/expansions/gather_spmd_expander.h | 2 ++ tensorflow/dtensor/mlir/expansions/identity_n_spmd_expander.cc | 2 ++ tensorflow/dtensor/mlir/expansions/io_op_spmd_expander.cc | 1 + tensorflow/dtensor/mlir/expansions/iterator_spmd_expander.cc | 1 + tensorflow/dtensor/mlir/expansions/matmul_spmd_expander.cc | 1 + tensorflow/dtensor/mlir/expansions/meta_spmd_expander.cc | 1 + tensorflow/dtensor/mlir/expansions/nullary_spmd_expander.cc | 2 ++ tensorflow/dtensor/mlir/expansions/reduce_spmd_expander.cc | 1 + tensorflow/dtensor/mlir/expansions/replicated_spmd_expander.cc | 1 + tensorflow/dtensor/mlir/expansions/scatter_spmd_expander.cc | 1 + .../dtensor/mlir/expansions/sparse_to_dense_spmd_expander.cc | 1 + tensorflow/dtensor/mlir/expansions/split_spmd_expander.cc | 1 + tensorflow/dtensor/mlir/expansions/squeeze_spmd_expander.cc | 1 + .../dtensor/mlir/expansions/tensorlist_setitem_spmd_expander.cc | 1 + tensorflow/dtensor/mlir/expansions/top_k_spmd_expander.cc | 1 + tensorflow/dtensor/mlir/expansions/trivial_spmd_expander.cc | 1 + .../dtensor/mlir/expansions/unsupported_op_spmd_expander.cc | 1 + 26 files changed, 31 insertions(+), 1 deletion(-) diff --git a/tensorflow/dtensor/mlir/expansions/argmax_spmd_expander.cc b/tensorflow/dtensor/mlir/expansions/argmax_spmd_expander.cc index 68c16e6035b419..729df319096ef7 100644 --- a/tensorflow/dtensor/mlir/expansions/argmax_spmd_expander.cc +++ b/tensorflow/dtensor/mlir/expansions/argmax_spmd_expander.cc @@ -19,6 +19,8 @@ limitations under the License. #include #include +#include "absl/status/status.h" +#include "absl/strings/str_cat.h" #include "llvm/ADT/DenseMap.h" #include "llvm/Support/Casting.h" #include "mlir/IR/Builders.h" // from @llvm-project diff --git a/tensorflow/dtensor/mlir/expansions/broadcast_to_spmd_expander.cc b/tensorflow/dtensor/mlir/expansions/broadcast_to_spmd_expander.cc index 5f7392aa80fb4e..9bf0f60226bf76 100644 --- a/tensorflow/dtensor/mlir/expansions/broadcast_to_spmd_expander.cc +++ b/tensorflow/dtensor/mlir/expansions/broadcast_to_spmd_expander.cc @@ -19,6 +19,7 @@ limitations under the License. #include #include +#include "absl/status/status.h" #include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/SmallPtrSet.h" diff --git a/tensorflow/dtensor/mlir/expansions/control_flow_spmd_expander.cc b/tensorflow/dtensor/mlir/expansions/control_flow_spmd_expander.cc index 127b1d1848e192..7df8dc12f654a2 100644 --- a/tensorflow/dtensor/mlir/expansions/control_flow_spmd_expander.cc +++ b/tensorflow/dtensor/mlir/expansions/control_flow_spmd_expander.cc @@ -17,6 +17,7 @@ limitations under the License. #include +#include "absl/status/status.h" #include "llvm/ADT/STLExtras.h" #include "llvm/Support/Casting.h" #include "mlir/IR/BuiltinTypes.h" // from @llvm-project diff --git a/tensorflow/dtensor/mlir/expansions/conv_spmd_expander.cc b/tensorflow/dtensor/mlir/expansions/conv_spmd_expander.cc index 4e069cf72b89c3..23052a2b5ac767 100644 --- a/tensorflow/dtensor/mlir/expansions/conv_spmd_expander.cc +++ b/tensorflow/dtensor/mlir/expansions/conv_spmd_expander.cc @@ -22,6 +22,7 @@ limitations under the License. #include #include "absl/status/status.h" +#include "absl/strings/str_cat.h" #include "llvm/ADT/SmallVector.h" #include "llvm/Support/Casting.h" #include "llvm/Support/FormatVariadic.h" diff --git a/tensorflow/dtensor/mlir/expansions/cumsum_spmd_expander.cc b/tensorflow/dtensor/mlir/expansions/cumsum_spmd_expander.cc index 9389ef58e6dd1f..28a0cd1ba5ff96 100644 --- a/tensorflow/dtensor/mlir/expansions/cumsum_spmd_expander.cc +++ b/tensorflow/dtensor/mlir/expansions/cumsum_spmd_expander.cc @@ -17,8 +17,8 @@ limitations under the License. #include #include -#include +#include "absl/status/status.h" #include "absl/strings/str_cat.h" #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/SmallPtrSet.h" diff --git a/tensorflow/dtensor/mlir/expansions/dataparallel_spmd_expander.cc b/tensorflow/dtensor/mlir/expansions/dataparallel_spmd_expander.cc index 9449e7540b00db..4378b1d9835904 100644 --- a/tensorflow/dtensor/mlir/expansions/dataparallel_spmd_expander.cc +++ b/tensorflow/dtensor/mlir/expansions/dataparallel_spmd_expander.cc @@ -20,6 +20,7 @@ limitations under the License. #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" +#include "absl/status/status.h" #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallPtrSet.h" diff --git a/tensorflow/dtensor/mlir/expansions/einsum_spmd_expander.cc b/tensorflow/dtensor/mlir/expansions/einsum_spmd_expander.cc index d49dfe5493d95c..d7e3d037fd9aed 100644 --- a/tensorflow/dtensor/mlir/expansions/einsum_spmd_expander.cc +++ b/tensorflow/dtensor/mlir/expansions/einsum_spmd_expander.cc @@ -24,6 +24,7 @@ limitations under the License. #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" #include "absl/status/status.h" +#include "absl/strings/str_cat.h" #include "absl/strings/str_split.h" #include "absl/strings/string_view.h" #include "llvm/ADT/DenseMap.h" diff --git a/tensorflow/dtensor/mlir/expansions/elementwise_spmd_expander.cc b/tensorflow/dtensor/mlir/expansions/elementwise_spmd_expander.cc index 9d0017c2a8bef3..2eda8eae30c16c 100644 --- a/tensorflow/dtensor/mlir/expansions/elementwise_spmd_expander.cc +++ b/tensorflow/dtensor/mlir/expansions/elementwise_spmd_expander.cc @@ -20,6 +20,7 @@ limitations under the License. #include #include "absl/container/flat_hash_set.h" +#include "absl/status/status.h" #include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/STLExtras.h" diff --git a/tensorflow/dtensor/mlir/expansions/fill_spmd_expander.cc b/tensorflow/dtensor/mlir/expansions/fill_spmd_expander.cc index 8bab4e55517dea..1d0b051bf0f6d0 100644 --- a/tensorflow/dtensor/mlir/expansions/fill_spmd_expander.cc +++ b/tensorflow/dtensor/mlir/expansions/fill_spmd_expander.cc @@ -18,6 +18,8 @@ limitations under the License. #include #include +#include "absl/status/status.h" +#include "absl/strings/str_cat.h" #include "llvm/Support/Casting.h" #include "mlir/IR/Block.h" // from @llvm-project #include "mlir/IR/Builders.h" // from @llvm-project diff --git a/tensorflow/dtensor/mlir/expansions/gather_spmd_expander.h b/tensorflow/dtensor/mlir/expansions/gather_spmd_expander.h index 7c4d67b7ce29a4..392d0cbf205caa 100644 --- a/tensorflow/dtensor/mlir/expansions/gather_spmd_expander.h +++ b/tensorflow/dtensor/mlir/expansions/gather_spmd_expander.h @@ -20,6 +20,8 @@ limitations under the License. #include #include +#include "absl/status/status.h" +#include "absl/strings/str_cat.h" #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/SmallVector.h" #include "llvm/Support/Casting.h" diff --git a/tensorflow/dtensor/mlir/expansions/identity_n_spmd_expander.cc b/tensorflow/dtensor/mlir/expansions/identity_n_spmd_expander.cc index af02f6a2fb1b7b..8ff3cae1a8f1ec 100644 --- a/tensorflow/dtensor/mlir/expansions/identity_n_spmd_expander.cc +++ b/tensorflow/dtensor/mlir/expansions/identity_n_spmd_expander.cc @@ -15,6 +15,8 @@ limitations under the License. #include "tensorflow/dtensor/mlir/expansions/identity_n_spmd_expander.h" +#include "absl/status/status.h" +#include "absl/strings/str_cat.h" #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/SmallPtrSet.h" #include "llvm/ADT/SmallVector.h" diff --git a/tensorflow/dtensor/mlir/expansions/io_op_spmd_expander.cc b/tensorflow/dtensor/mlir/expansions/io_op_spmd_expander.cc index 4620c9fdcaa0dd..b8b98c651591c5 100644 --- a/tensorflow/dtensor/mlir/expansions/io_op_spmd_expander.cc +++ b/tensorflow/dtensor/mlir/expansions/io_op_spmd_expander.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include "absl/status/status.h" #include "llvm/Support/Casting.h" #include "llvm/Support/FormatVariadic.h" #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project diff --git a/tensorflow/dtensor/mlir/expansions/iterator_spmd_expander.cc b/tensorflow/dtensor/mlir/expansions/iterator_spmd_expander.cc index 7236d5a30c80f4..75c3e893cf2019 100644 --- a/tensorflow/dtensor/mlir/expansions/iterator_spmd_expander.cc +++ b/tensorflow/dtensor/mlir/expansions/iterator_spmd_expander.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include "absl/status/status.h" #include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/SmallVector.h" #include "llvm/Support/FormatVariadic.h" diff --git a/tensorflow/dtensor/mlir/expansions/matmul_spmd_expander.cc b/tensorflow/dtensor/mlir/expansions/matmul_spmd_expander.cc index d1a56c89ae9905..3f1d3e136464da 100644 --- a/tensorflow/dtensor/mlir/expansions/matmul_spmd_expander.cc +++ b/tensorflow/dtensor/mlir/expansions/matmul_spmd_expander.cc @@ -24,6 +24,7 @@ limitations under the License. #include "absl/container/flat_hash_set.h" #include "absl/status/status.h" +#include "absl/strings/str_cat.h" #include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/SmallSet.h" diff --git a/tensorflow/dtensor/mlir/expansions/meta_spmd_expander.cc b/tensorflow/dtensor/mlir/expansions/meta_spmd_expander.cc index 74fc87afce53f6..2b2b917dbb11f5 100644 --- a/tensorflow/dtensor/mlir/expansions/meta_spmd_expander.cc +++ b/tensorflow/dtensor/mlir/expansions/meta_spmd_expander.cc @@ -23,6 +23,7 @@ limitations under the License. #include #include "absl/status/status.h" +#include "absl/strings/str_cat.h" #include "absl/types/optional.h" #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/STLExtras.h" diff --git a/tensorflow/dtensor/mlir/expansions/nullary_spmd_expander.cc b/tensorflow/dtensor/mlir/expansions/nullary_spmd_expander.cc index 4053cc70ecc028..e6f79b1880ac60 100644 --- a/tensorflow/dtensor/mlir/expansions/nullary_spmd_expander.cc +++ b/tensorflow/dtensor/mlir/expansions/nullary_spmd_expander.cc @@ -19,6 +19,8 @@ limitations under the License. #include #include +#include "absl/status/status.h" +#include "absl/strings/str_cat.h" #include "llvm/ADT/SmallPtrSet.h" #include "llvm/ADT/SmallVector.h" #include "llvm/Support/FormatVariadic.h" diff --git a/tensorflow/dtensor/mlir/expansions/reduce_spmd_expander.cc b/tensorflow/dtensor/mlir/expansions/reduce_spmd_expander.cc index cafcfe0c38ad5f..8714e6ba585cac 100644 --- a/tensorflow/dtensor/mlir/expansions/reduce_spmd_expander.cc +++ b/tensorflow/dtensor/mlir/expansions/reduce_spmd_expander.cc @@ -21,6 +21,7 @@ limitations under the License. #include "absl/container/flat_hash_set.h" #include "absl/status/status.h" +#include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/SmallPtrSet.h" diff --git a/tensorflow/dtensor/mlir/expansions/replicated_spmd_expander.cc b/tensorflow/dtensor/mlir/expansions/replicated_spmd_expander.cc index 477506e5a5366d..2cb720c2944ccb 100644 --- a/tensorflow/dtensor/mlir/expansions/replicated_spmd_expander.cc +++ b/tensorflow/dtensor/mlir/expansions/replicated_spmd_expander.cc @@ -17,6 +17,7 @@ limitations under the License. #include +#include "absl/status/status.h" #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/SmallPtrSet.h" #include "llvm/ADT/SmallVector.h" diff --git a/tensorflow/dtensor/mlir/expansions/scatter_spmd_expander.cc b/tensorflow/dtensor/mlir/expansions/scatter_spmd_expander.cc index bd5612f20781fa..c99c6492342ba3 100644 --- a/tensorflow/dtensor/mlir/expansions/scatter_spmd_expander.cc +++ b/tensorflow/dtensor/mlir/expansions/scatter_spmd_expander.cc @@ -21,6 +21,7 @@ limitations under the License. #include #include "absl/container/flat_hash_set.h" +#include "absl/status/status.h" #include "absl/strings/str_cat.h" #include "absl/types/optional.h" #include "llvm/ADT/ArrayRef.h" diff --git a/tensorflow/dtensor/mlir/expansions/sparse_to_dense_spmd_expander.cc b/tensorflow/dtensor/mlir/expansions/sparse_to_dense_spmd_expander.cc index 7e1496c6c1cc79..e820bc847ec33f 100644 --- a/tensorflow/dtensor/mlir/expansions/sparse_to_dense_spmd_expander.cc +++ b/tensorflow/dtensor/mlir/expansions/sparse_to_dense_spmd_expander.cc @@ -17,6 +17,7 @@ limitations under the License. #include +#include "absl/status/status.h" #include "mlir/IR/BuiltinTypes.h" // from @llvm-project #include "mlir/IR/Operation.h" // from @llvm-project #include "mlir/Support/LLVM.h" // from @llvm-project diff --git a/tensorflow/dtensor/mlir/expansions/split_spmd_expander.cc b/tensorflow/dtensor/mlir/expansions/split_spmd_expander.cc index ac183c765f8872..dea969931001d1 100644 --- a/tensorflow/dtensor/mlir/expansions/split_spmd_expander.cc +++ b/tensorflow/dtensor/mlir/expansions/split_spmd_expander.cc @@ -20,6 +20,7 @@ limitations under the License. #include #include +#include "absl/status/status.h" #include "llvm/ADT/DenseMap.h" #include "mlir/IR/Value.h" // from @llvm-project #include "mlir/Support/LLVM.h" // from @llvm-project diff --git a/tensorflow/dtensor/mlir/expansions/squeeze_spmd_expander.cc b/tensorflow/dtensor/mlir/expansions/squeeze_spmd_expander.cc index 80bbf6967b38fd..78b14dbc3b87ff 100644 --- a/tensorflow/dtensor/mlir/expansions/squeeze_spmd_expander.cc +++ b/tensorflow/dtensor/mlir/expansions/squeeze_spmd_expander.cc @@ -21,6 +21,7 @@ limitations under the License. #include #include +#include "absl/status/status.h" #include "mlir/IR/Builders.h" // from @llvm-project #include "mlir/IR/BuiltinAttributes.h" // from @llvm-project #include "mlir/Support/LLVM.h" // from @llvm-project diff --git a/tensorflow/dtensor/mlir/expansions/tensorlist_setitem_spmd_expander.cc b/tensorflow/dtensor/mlir/expansions/tensorlist_setitem_spmd_expander.cc index 30605ee5f305d7..451bf2acc64413 100644 --- a/tensorflow/dtensor/mlir/expansions/tensorlist_setitem_spmd_expander.cc +++ b/tensorflow/dtensor/mlir/expansions/tensorlist_setitem_spmd_expander.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/dtensor/mlir/expansions/tensorlist_setitem_spmd_expander.h" +#include "absl/status/status.h" #include "llvm/ADT/DenseMap.h" #include "mlir/IR/Operation.h" // from @llvm-project #include "tensorflow/core/platform/errors.h" diff --git a/tensorflow/dtensor/mlir/expansions/top_k_spmd_expander.cc b/tensorflow/dtensor/mlir/expansions/top_k_spmd_expander.cc index cb14920d7018f4..c36da3e5f93b70 100644 --- a/tensorflow/dtensor/mlir/expansions/top_k_spmd_expander.cc +++ b/tensorflow/dtensor/mlir/expansions/top_k_spmd_expander.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include "absl/status/status.h" #include "llvm/ADT/DenseMap.h" #include "mlir/IR/Builders.h" // from @llvm-project #include "mlir/IR/IRMapping.h" // from @llvm-project diff --git a/tensorflow/dtensor/mlir/expansions/trivial_spmd_expander.cc b/tensorflow/dtensor/mlir/expansions/trivial_spmd_expander.cc index d881366c462efa..0fab24b6b498a8 100644 --- a/tensorflow/dtensor/mlir/expansions/trivial_spmd_expander.cc +++ b/tensorflow/dtensor/mlir/expansions/trivial_spmd_expander.cc @@ -17,6 +17,7 @@ limitations under the License. #include +#include "absl/status/status.h" #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallVector.h" diff --git a/tensorflow/dtensor/mlir/expansions/unsupported_op_spmd_expander.cc b/tensorflow/dtensor/mlir/expansions/unsupported_op_spmd_expander.cc index 85441584b642ca..8e68266ec9e9e5 100644 --- a/tensorflow/dtensor/mlir/expansions/unsupported_op_spmd_expander.cc +++ b/tensorflow/dtensor/mlir/expansions/unsupported_op_spmd_expander.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/dtensor/mlir/expansions/unsupported_op_spmd_expander.h" +#include "absl/status/status.h" #include "absl/strings/string_view.h" #include "llvm/ADT/DenseMap.h" #include "mlir/IR/Operation.h" // from @llvm-project From da79958e37365abd0471c6feaa4ff4364629f368 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Fri, 28 Aug 2026 23:09:26 -0700 Subject: [PATCH 11/12] Automated Code Change PiperOrigin-RevId: 972978443 --- third_party/xla/xla/pjrt/se/BUILD | 1 + third_party/xla/xla/pjrt/se/local_device_state.cc | 1 - third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc | 1 - third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc | 1 + third_party/xla/xla/pjrt/se/stream_executor_executable.cc | 2 +- third_party/xla/xla/pjrt/se/stream_executor_executable.h | 1 + third_party/xla/xla/pjrt/se/stream_executor_executable_test.cc | 1 - 7 files changed, 4 insertions(+), 4 deletions(-) diff --git a/third_party/xla/xla/pjrt/se/BUILD b/third_party/xla/xla/pjrt/se/BUILD index 61eda90fc9c6cc..1fa90239879983 100644 --- a/third_party/xla/xla/pjrt/se/BUILD +++ b/third_party/xla/xla/pjrt/se/BUILD @@ -363,6 +363,7 @@ xla_cc_test( "//xla/tsl/platform:env", "//xla/tsl/platform:statusor", "@com_google_absl//absl/functional:any_invocable", + "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_matchers", diff --git a/third_party/xla/xla/pjrt/se/local_device_state.cc b/third_party/xla/xla/pjrt/se/local_device_state.cc index 11b0fab30eab7c..13ab5fee8c0a2c 100644 --- a/third_party/xla/xla/pjrt/se/local_device_state.cc +++ b/third_party/xla/xla/pjrt/se/local_device_state.cc @@ -17,7 +17,6 @@ limitations under the License. #include #include -#include #include #include #include diff --git a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc index 870d2482bdb027..e680c677ef0429 100644 --- a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc +++ b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc @@ -71,7 +71,6 @@ limitations under the License. #include #include #include -#include #include #include #include diff --git a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc index e34af9ccc21186..7449168d3d26a6 100644 --- a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc +++ b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc @@ -29,6 +29,7 @@ limitations under the License. #include #include "absl/functional/any_invocable.h" #include "absl/log/check.h" +#include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/status_matchers.h" #include "absl/status/statusor.h" diff --git a/third_party/xla/xla/pjrt/se/stream_executor_executable.cc b/third_party/xla/xla/pjrt/se/stream_executor_executable.cc index 94eaddff745c85..76f57ba3745a30 100644 --- a/third_party/xla/xla/pjrt/se/stream_executor_executable.cc +++ b/third_party/xla/xla/pjrt/se/stream_executor_executable.cc @@ -16,7 +16,6 @@ limitations under the License. #include "xla/pjrt/se/stream_executor_executable.h" #include -#include #include #include #include @@ -30,6 +29,7 @@ limitations under the License. #include "absl/status/status_macros.h" #include "absl/strings/cord.h" #include "absl/strings/string_view.h" +#include "absl/synchronization/mutex.h" #include "riegeli/base/any.h" #include "riegeli/base/maker.h" #include "riegeli/bytes/cord_reader.h" diff --git a/third_party/xla/xla/pjrt/se/stream_executor_executable.h b/third_party/xla/xla/pjrt/se/stream_executor_executable.h index 6fac5f281fd9d4..ad5b1225b63ec8 100644 --- a/third_party/xla/xla/pjrt/se/stream_executor_executable.h +++ b/third_party/xla/xla/pjrt/se/stream_executor_executable.h @@ -39,6 +39,7 @@ limitations under the License. #include "xla/pjrt/pjrt_common.h" #include "xla/pjrt/pjrt_compiler.h" #include "xla/pjrt/pjrt_executable.h" +#include "xla/pjrt/proto/compile_options.pb.h" #include "xla/service/compiled_module.h" #include "xla/service/hlo.pb.h" #include "xla/stream_executor/abi/executable_abi_version.h" diff --git a/third_party/xla/xla/pjrt/se/stream_executor_executable_test.cc b/third_party/xla/xla/pjrt/se/stream_executor_executable_test.cc index 2da6cb9a4bab5b..c90d81f68bca8c 100644 --- a/third_party/xla/xla/pjrt/se/stream_executor_executable_test.cc +++ b/third_party/xla/xla/pjrt/se/stream_executor_executable_test.cc @@ -17,7 +17,6 @@ limitations under the License. #include #include -#include #include #include From 1397bef384b069c7375133194c7b843dda6330d5 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Fri, 28 Aug 2026 23:23:13 -0700 Subject: [PATCH 12/12] Automated Code Change PiperOrigin-RevId: 972982809 --- third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc | 14 ++++++-------- .../pjrt/gpu/se_gpu_pjrt_client_multi_gpu_test.cc | 8 ++++---- 2 files changed, 10 insertions(+), 12 deletions(-) diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc index f9680e7caad4c8..27333f76cdc530 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc @@ -296,8 +296,7 @@ void StreamExecutorGpuRawClient::UpdateGlobalProcessInfo( void StreamExecutorGpuRawClient::UpdateCompileOptionsTopology( const PjRtTopologyDescription& topology, CompileOptions* options) const { options->executable_build_options.set_gpu_topology( - tensorflow::down_cast( - &topology) + absl::down_cast(&topology) ->gpu_topology()); } @@ -310,7 +309,7 @@ namespace { // Get the local device state for a given PjRtDevice. absl::StatusOr GetLocalDeviceState(PjRtDevice* device) { PjRtStreamExecutorDevice* pjrt_se_device = - tensorflow::down_cast(device); + absl::down_cast(device); return pjrt_se_device->GetLocalDeviceState(); } @@ -629,7 +628,7 @@ StreamExecutorGpuRawClient::CrossHostTransferBuffers( // Get the local_device_state and use it to schedule transfers. Fail // transfers early if we cannot get the local_device_state. absl::StatusOr local_device_state = - tensorflow::down_cast(device) + absl::down_cast(device) ->GetLocalDeviceState(); if (!local_device_state.ok()) { SetEventAsError(transfer_event, local_device_state.status()); @@ -1833,7 +1832,7 @@ absl::StatusOr StreamExecutorGpuDevice::GetAllocatorStats() } auto* allocator_adapter = dynamic_cast( - tensorflow::down_cast(client())->allocator()); + absl::down_cast(client())->allocator()); if (!allocator_adapter) { return Unimplemented( "GetAllocatorStats() is only implemented with MultiDeviceAdapter " @@ -1858,7 +1857,7 @@ absl::Status StreamExecutorGpuDevice::ClearMemoryStats() { } auto* allocator_adapter = dynamic_cast( - tensorflow::down_cast(client())->allocator()); + absl::down_cast(client())->allocator()); if (!allocator_adapter) { return absl::UnimplementedError( "ClearMemoryStats() is only implemented with MultiDeviceAdapter " @@ -2224,8 +2223,7 @@ static absl::StatusOr RunGpuAsync( ABSL_ASSIGN_OR_RETURN(auto options_and_stream, exec.RunHelper(argument_shapes, run_options_inp)); - auto* gpu_exec = - tensorflow::down_cast(exec.executable()); + auto* gpu_exec = absl::down_cast(exec.executable()); const ServiceExecutableRunOptions* run_options = &options_and_stream.first; se::DeviceAddressAllocator* const memory_allocator = run_options->allocator(); diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_multi_gpu_test.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_multi_gpu_test.cc index c390eda8c6139f..2801598a403a7a 100644 --- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_multi_gpu_test.cc +++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_multi_gpu_test.cc @@ -313,7 +313,7 @@ TEST(StreamExecutorGpuClientTest, } std::unique_ptr& client = *client_status; auto* gpu_client = - tsl::down_cast(client.get()); + absl::down_cast(client.get()); const gpu::GpuExecutableRunOptions* run_options = gpu_client->gpu_run_options(); if (run_options == nullptr || @@ -345,7 +345,7 @@ TEST(StreamExecutorGpuClientTest, options.abort_collectives_on_failure = true; ASSERT_OK_AND_ASSIGN(auto client, GetStreamExecutorGpuClient(options)); - auto* gpu_client = tsl::down_cast(client.get()); + auto* gpu_client = absl::down_cast(client.get()); const gpu::GpuExecutableRunOptions* run_options = gpu_client->gpu_run_options(); ASSERT_NE(run_options, nullptr); @@ -415,7 +415,7 @@ TEST(StreamExecutorGpuClientTest, } auto* gpu_client0 = - tsl::down_cast(pjrt_clients[0].get()); + absl::down_cast(pjrt_clients[0].get()); const gpu::GpuExecutableRunOptions* run_options = gpu_client0->gpu_run_options(); ASSERT_NE(run_options, nullptr); @@ -1676,7 +1676,7 @@ absl::Status InterProcessCollectiveInitTestBody(int rank_id) { // executor's collective memory allocator into the selected collectives // backend (e.g. MORI ShmemMalloc). With inert backend stubs the allocation // may return null; we only log the outcome and do not fail the test. - auto* se_device = tsl::down_cast( + auto* se_device = absl::down_cast( client->addressable_devices()[0]); TF_RET_CHECK(se_device != nullptr); LocalDeviceState* local_device_state = se_device->local_device_state();