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 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 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 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. 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 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); }); } 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) 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(); 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 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/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; 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( 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