diff --git a/tensorflow/compiler/mlir/lite/core/api/flatbuffer_conversions.cc b/tensorflow/compiler/mlir/lite/core/api/flatbuffer_conversions.cc index aace2080a0481f..2c4b62b39ec179 100644 --- a/tensorflow/compiler/mlir/lite/core/api/flatbuffer_conversions.cc +++ b/tensorflow/compiler/mlir/lite/core/api/flatbuffer_conversions.cc @@ -406,12 +406,12 @@ void CheckParsePointerParams(const Operator* op, TFLITE_DCHECK(builtin_data != nullptr); } -// Copies the contents from the flatbuffer int vector `flatbuffer` into the +// Copies the contents from the flatbuffer int vector `flat_vector` into the // int array `buffer`. `flat_vector` and `buffer` represent the same // configuration operation for a given operation. template static absl::Status FlatBufferIntVectorToArray( - int max_size_of_buffer, const flatbuffers::Vector* flat_vector, + size_t max_size_of_buffer, const flatbuffers::Vector* flat_vector, DataType* buffer, const char* op_name) { if (!flat_vector) { auto error_message = absl::StrFormat( @@ -2498,7 +2498,7 @@ absl::Status ParseStablehloScatter(const Operator* op, if (schema_params->update_window_dims()) { TFL_FILE_ENSURE_STATUS(FlatBufferIntVectorToArray( - schema_params->update_window_dims()->size() * sizeof(int64_t), + sizeof(params->update_window_dims), schema_params->update_window_dims(), params->update_window_dims, "stablehlo_scatter")); params->num_update_window_dims = @@ -2507,7 +2507,7 @@ absl::Status ParseStablehloScatter(const Operator* op, if (schema_params->inserted_window_dims()) { TFL_FILE_ENSURE_STATUS(FlatBufferIntVectorToArray( - schema_params->inserted_window_dims()->size() * sizeof(int64_t), + sizeof(params->inserted_window_dims), schema_params->inserted_window_dims(), params->inserted_window_dims, "stablehlo_scatter")); params->num_inserted_window_dims = @@ -2516,8 +2516,7 @@ absl::Status ParseStablehloScatter(const Operator* op, if (schema_params->scatter_dims_to_operand_dims()) { TFL_FILE_ENSURE_STATUS(FlatBufferIntVectorToArray( - schema_params->scatter_dims_to_operand_dims()->size() * - sizeof(int64_t), + sizeof(params->scatter_dims_to_operand_dims), schema_params->scatter_dims_to_operand_dims(), params->scatter_dims_to_operand_dims, "stablehlo_scatter")); params->num_scatter_dims_to_operand_dims = @@ -2579,8 +2578,7 @@ absl::Status ParseStablehloGather(const Operator* op, if (schema_params != nullptr) { if (schema_params->offset_dims()) { TFL_FILE_ENSURE_STATUS(FlatBufferIntVectorToArray( - /*max_size_of_buffer=*/schema_params->offset_dims()->size() * - sizeof(int64_t), + /*max_size_of_buffer=*/sizeof(params->offset_dims), /*flat_vector=*/schema_params->offset_dims(), /*buffer=*/params->offset_dims, /*op_name=*/"stablehlo_gather")); @@ -2589,7 +2587,7 @@ absl::Status ParseStablehloGather(const Operator* op, if (schema_params->collapsed_slice_dims()) { TFL_FILE_ENSURE_STATUS(FlatBufferIntVectorToArray( - schema_params->collapsed_slice_dims()->size() * sizeof(int64_t), + sizeof(params->collapsed_slice_dims), schema_params->collapsed_slice_dims(), params->collapsed_slice_dims, "stablehlo_gather")); params->num_collapsed_slice_dims = @@ -2598,9 +2596,8 @@ absl::Status ParseStablehloGather(const Operator* op, if (schema_params->start_index_map()) { TFL_FILE_ENSURE_STATUS(FlatBufferIntVectorToArray( - schema_params->start_index_map()->size() * sizeof(int64_t), - schema_params->start_index_map(), params->start_index_map, - "stablehlo_gather")); + sizeof(params->start_index_map), schema_params->start_index_map(), + params->start_index_map, "stablehlo_gather")); params->num_start_index_map = schema_params->start_index_map()->size(); } @@ -2608,9 +2605,8 @@ absl::Status ParseStablehloGather(const Operator* op, if (schema_params->slice_sizes()) { TFL_FILE_ENSURE_STATUS(FlatBufferIntVectorToArray( - schema_params->slice_sizes()->size() * sizeof(int64_t), - schema_params->slice_sizes(), params->slice_sizes, - "stablehlo_gather")); + sizeof(params->slice_sizes), schema_params->slice_sizes(), + params->slice_sizes, "stablehlo_gather")); params->num_slice_sizes = schema_params->slice_sizes()->size(); } diff --git a/tensorflow/compiler/mlir/lite/core/api/flatbuffer_conversions_test.cc b/tensorflow/compiler/mlir/lite/core/api/flatbuffer_conversions_test.cc index 7bfd7fa15df945..cd6f17154219d2 100644 --- a/tensorflow/compiler/mlir/lite/core/api/flatbuffer_conversions_test.cc +++ b/tensorflow/compiler/mlir/lite/core/api/flatbuffer_conversions_test.cc @@ -49,18 +49,24 @@ using tflite::BuiltinOperator_CUSTOM; using tflite::BuiltinOperator_FULLY_CONNECTED; using tflite::BuiltinOperator_RESHAPE; using tflite::BuiltinOperator_SQUEEZE; +using tflite::BuiltinOperator_STABLEHLO_GATHER; using tflite::BuiltinOperator_STABLEHLO_PAD; using tflite::BuiltinOperator_STABLEHLO_REDUCE_WINDOW; +using tflite::BuiltinOperator_STABLEHLO_SCATTER; +using tflite::BuiltinOptions2_StablehloGatherOptions; using tflite::BuiltinOptions2_StablehloPadOptions; using tflite::BuiltinOptions2_StablehloReduceWindowOptions; +using tflite::BuiltinOptions2_StablehloScatterOptions; using tflite::BuiltinOptions_Conv2DOptions; using tflite::BuiltinOptions_FullyConnectedOptions; using tflite::BuiltinOptions_NONE; using tflite::BuiltinOptions_ReshapeOptions; using tflite::CreateReshapeOptions; using tflite::CreateSqueezeOptions; +using tflite::CreateStablehloGatherOptions; using tflite::CreateStablehloPadOptions; using tflite::CreateStablehloReduceWindowOptions; +using tflite::CreateStablehloScatterOptions; using tflite::FullyConnectedOptionsWeightsFormat; using tflite::Padding_SAME; using tflite::TensorType_BFLOAT16; @@ -871,5 +877,99 @@ TEST_F(StablehloPadFlatbufferConversionsTest, DeathTests) { ""); } +class StablehloScatterFlatbufferConversionsTest + : public FlatbufferConversionsTest { + protected: + static constexpr int kMaxDims = + TFLITE_STABLEHLO_SCATTER_PARAMS_MAX_DIMENSION_COUNT; + + const Operator* BuildScatterOperator(int update_window_dims_count, + int inserted_window_dims_count, + int scatter_dims_count) { + const auto update_window_dims = builder_.CreateVector( + std::vector(update_window_dims_count, 1)); + const auto inserted_window_dims = builder_.CreateVector( + std::vector(inserted_window_dims_count, 1)); + const auto scatter_dims = + builder_.CreateVector(std::vector(scatter_dims_count, 1)); + return BuildTestOperator( + BuiltinOptions2_StablehloScatterOptions, + CreateStablehloScatterOptions( + builder_, /*indices_are_sorted=*/true, update_window_dims, + inserted_window_dims, scatter_dims, + /*index_vector_dim=*/0, /*unique_indices=*/false, + /*update_computation_subgraph_index=*/0) + .Union()); + } + + void ExpectTooManyDimensions(int update_window_dims_count, + int inserted_window_dims_count, + int scatter_dims_count) { + TfLiteStablehloScatterParams* output_data = nullptr; + const auto status = ParseOpData( + BuildScatterOperator(update_window_dims_count, + inserted_window_dims_count, scatter_dims_count), + BuiltinOperator_STABLEHLO_SCATTER, &mock_allocator_, + (void**)&output_data); + EXPECT_FALSE(status.ok()); + EXPECT_EQ(output_data, nullptr); + EXPECT_THAT(status.message(), + HasSubstr("Found too many dimensions in the input array of " + "operation 'stablehlo_scatter'.")); + } +}; + +TEST_F(StablehloScatterFlatbufferConversionsTest, AcceptsMaximumDimensions) { + TfLiteStablehloScatterParams* output_data = nullptr; + const auto status = + ParseOpData(BuildScatterOperator(kMaxDims, kMaxDims, kMaxDims), + BuiltinOperator_STABLEHLO_SCATTER, &mock_allocator_, + (void**)&output_data); + EXPECT_TRUE(status.ok()); + ASSERT_NE(output_data, nullptr); + EXPECT_EQ(output_data->num_update_window_dims, kMaxDims); + EXPECT_EQ(output_data->num_inserted_window_dims, kMaxDims); + EXPECT_EQ(output_data->num_scatter_dims_to_operand_dims, kMaxDims); +} + +TEST_F(StablehloScatterFlatbufferConversionsTest, + RejectsTooManyUpdateWindowDimensions) { + ExpectTooManyDimensions(kMaxDims + 1, 1, 1); +} + +TEST_F(StablehloScatterFlatbufferConversionsTest, + RejectsTooManyInsertedWindowDimensions) { + ExpectTooManyDimensions(1, kMaxDims + 1, 1); +} + +TEST_F(StablehloScatterFlatbufferConversionsTest, + RejectsTooManyScatterDimensions) { + ExpectTooManyDimensions(1, 1, kMaxDims + 1); +} + +TEST_F(FlatbufferConversionsTest, + ParseStablehloGatherRejectsTooManyDimensions) { + std::vector too_many_dims( + TFLITE_STABLEHLO_GATHER_PARAMS_MAX_DIMENSION_COUNT + 1, 1); + const Operator* op = BuildTestOperator( + BuiltinOptions2_StablehloGatherOptions, + CreateStablehloGatherOptions( + builder_, /*offset_dims=*/builder_.CreateVector(too_many_dims), + /*collapsed_slice_dims=*/builder_.CreateVector({1}), + /*start_index_map=*/builder_.CreateVector({1}), + /*index_vector_dim=*/0, + /*slice_sizes=*/builder_.CreateVector({1}), + /*indices_are_sorted=*/true) + .Union()); + void* output_data = nullptr; + const auto status = ParseOpData(op, BuiltinOperator_STABLEHLO_GATHER, + &mock_allocator_, &output_data); + EXPECT_FALSE(status.ok()); + EXPECT_EQ(output_data, nullptr); + EXPECT_THAT(status.message(), + HasSubstr("Found too many dimensions in the input array of " + "operation 'stablehlo_gather'.")); +} + } // namespace flatbuffer_conversions } // namespace tflite_file diff --git a/tensorflow/compiler/mlir/lite/python/BUILD b/tensorflow/compiler/mlir/lite/python/BUILD index 27db8de43e7089..44b2ef7416e28e 100644 --- a/tensorflow/compiler/mlir/lite/python/BUILD +++ b/tensorflow/compiler/mlir/lite/python/BUILD @@ -91,7 +91,9 @@ cc_library( "//tensorflow/compiler/mlir/lite:tensorflow_lite", "//tensorflow/compiler/mlir/lite/core:macros", "//tensorflow/compiler/mlir/lite/debug", + "//tensorflow/compiler/mlir/lite/quantization/ir:QuantOps", "//tensorflow/compiler/mlir/lite/stablehlo:drop_shape_assertions", + "//tensorflow/compiler/mlir/lite/stablehlo:legalize_vhlo_quant_custom_calls", "//tensorflow/compiler/mlir/lite/stablehlo:prepare_hlo", "//tensorflow/compiler/mlir/lite/stablehlo:tfl_legalize_hlo", "//tensorflow/compiler/mlir/lite/stablehlo:unfold_splat_constant_pass", @@ -103,6 +105,8 @@ cc_library( "@llvm-project//mlir:FuncExtensions", "@llvm-project//mlir:IR", "@llvm-project//mlir:Pass", + "@llvm-project//mlir:QuantOps", + "@llvm-project//mlir:ReconcileUnrealizedCasts", "@llvm-project//mlir:Support", "@llvm-project//mlir:Transforms", "@stablehlo//:stablehlo_ops", diff --git a/tensorflow/compiler/mlir/lite/python/stablehlo_tfl_pipeline.cc b/tensorflow/compiler/mlir/lite/python/stablehlo_tfl_pipeline.cc index a48f944c2de7c9..59504630a1db05 100644 --- a/tensorflow/compiler/mlir/lite/python/stablehlo_tfl_pipeline.cc +++ b/tensorflow/compiler/mlir/lite/python/stablehlo_tfl_pipeline.cc @@ -23,8 +23,10 @@ limitations under the License. #include "absl/strings/str_cat.h" #include "llvm/Support/raw_ostream.h" #include "mlir/Bytecode/BytecodeWriter.h" // from @llvm-project +#include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h" // from @llvm-project #include "mlir/Dialect/Func/Extensions/InlinerExtension.h" // from @llvm-project #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project +#include "mlir/Dialect/Quant/IR/Quant.h" // from @llvm-project #include "mlir/IR/BuiltinOps.h" // from @llvm-project #include "mlir/IR/OperationSupport.h" // from @llvm-project #include "mlir/Pass/PassInstrumentation.h" // from @llvm-project @@ -40,6 +42,7 @@ limitations under the License. #include "tensorflow/compiler/mlir/lite/debug/debug.h" #include "tensorflow/compiler/mlir/lite/flatbuffer_export.h" #include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h" +#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h" #include "tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.h" #include "tensorflow/compiler/mlir/lite/transforms/cast_bf16_ops_to_f32_pass.h" #include "tensorflow/compiler/mlir/lite/transforms/large_constant_fold_pass.h" @@ -56,7 +59,7 @@ namespace mlir::TFL { void AddSkipToTflitePasses(mlir::OpPassManager& pass_manager) { pass_manager.addNestedPass( mlir::odml::CreateLegalizeChloToTflPass()); - pass_manager.addPass(mlir::odml::CreateCompositeLoweringPass()); + pass_manager.addPass(mlir::createInlinerPass()); pass_manager.addPass(mlir::TFL::CreateLowerQuantAnnotationsPass()); pass_manager.addPass(mlir::createSymbolDCEPass()); } @@ -65,6 +68,10 @@ void AddHloOptimizationPasses(mlir::OpPassManager& pass_manager) { // Drop shape assertion custom calls before VHLO legalization pass_manager.addPass(mlir::odml::CreateDropShapeAssertionsPass()); + // Legalize VHLO quant custom calls to StableHLO custom calls before VHLO + // legalization + pass_manager.addPass(mlir::odml::CreateLegalizeVhloQuantCustomCallsPass()); + // VHLO -> StableHLO pass_manager.addPass(mlir::stablehlo::createVhloLegalizeToStablehloPass()); @@ -108,6 +115,10 @@ void AddHloOptimizationPasses(mlir::OpPassManager& pass_manager) { // StableHLO -> MHLO bridge pass_manager.addPass(mlir::mhlo::createStablehloLegalizeToHloPass()); + // Composite lowering when IR is in MHLO + pass_manager.addPass(mlir::odml::CreateCompositeLoweringPass()); + pass_manager.addPass(mlir::createSymbolDCEPass()); + // MHLO algebraic optimizations pass_manager.addNestedPass( mlir::mhlo::createLegalizeEinsumToDotGeneralPass()); @@ -128,7 +139,13 @@ void AddHloToTfLiteLegalizationPasses(mlir::OpPassManager& pass_manager) { mlir::odml::CreateUniformQuantizedStableHloToTflPass()); pass_manager.addNestedPass( mlir::odml::CreatePrepareHloPass()); + pass_manager.addPass(mlir::odml::CreateUnfoldSplatConstantPass()); pass_manager.addPass(mlir::odml::CreateLegalizeHloToTfLitePass()); + + // Legalize remaining MHLO ops to StableHLO + pass_manager.addPass(mlir::mhlo::createHloLegalizeToStablehloPass()); + pass_manager.addNestedPass( + mlir::odml::createLegalizeCompositeToCustomOpPass()); } void AddTfLiteOptimizationPasses(mlir::OpPassManager& pass_manager, @@ -148,19 +165,23 @@ void AddTfLiteOptimizationPasses(mlir::OpPassManager& pass_manager, pass_manager.addNestedPass( mlir::TFL::CreateOptimizePass()); - // Quantization - pass_manager.addNestedPass( - mlir::TFL::CreatePrepareQuantizePass(pass_config.quant_specs)); - pass_manager.addNestedPass( - mlir::TFL::CreateQuantizePass(pass_config.quant_specs)); - pass_manager.addNestedPass( - mlir::TFL::CreatePostQuantizePass(/*emit_quant_adaptor_ops=*/true)); - if (!pass_config.unfold_batch_matmul) { pass_manager.addNestedPass( mlir::TFL::CreateOptimizeBatchMatmulPass()); + pass_manager.addNestedPass( + mlir::TFL::CreateOptimizePass()); } + // Quantization + pass_manager.addPass(mlir::TFL::CreatePropagateQParamsPass()); + pass_manager.addPass(mlir::TFL::CreateBiasQuantizerPass()); + pass_manager.addPass(mlir::TFL::CreateFuseQDQPass()); + + pass_manager.addNestedPass( + mlir::TFL::CreatePostQuantizePass(/*emit_quant_adaptor_ops=*/true)); + pass_manager.addNestedPass( + mlir::createCanonicalizerPass()); + // Some optimizations need to happen on the quantized graph. pass_manager.addNestedPass( mlir::TFL::CreateOptimizePass()); @@ -176,6 +197,8 @@ void AddTfLiteOptimizationPasses(mlir::OpPassManager& pass_manager, pass_manager.addNestedPass( mlir::createCanonicalizerPass()); pass_manager.addNestedPass(mlir::createCSEPass()); + pass_manager.addPass(mlir::TFL::CreateCleanupOptimizationBarrierPass()); + pass_manager.addPass(mlir::createReconcileUnrealizedCastsPass()); } absl::Status ConvertStableHloToTFLite( @@ -185,8 +208,10 @@ absl::Status ConvertStableHloToTFLite( mlir::MLIRContext* context = module->getContext(); mlir::DialectRegistry registry; mlir::func::registerInlinerExtension(registry); - registry.insert(); + registry.insert(); context->appendDialectRegistry(registry); context->loadAllAvailableDialects(); diff --git a/tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.cc b/tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.cc index fa86a8486913a0..a132320460fdce 100644 --- a/tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.cc +++ b/tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.cc @@ -334,18 +334,27 @@ TypeAttr RescaleQuantizedType(const Type input, const Attribute factor) { const auto element_type = quant::QuantizedType::getQuantizedElementType(input); if (!element_type) return {}; - if (auto qtype = dyn_cast(element_type)) { + if (const auto qtype = + dyn_cast(element_type)) { const ArrayRef scales = qtype.getScales(); - // Broadcasting hasn't been implemented yet. - if (static_cast(scales.size()) != factor_values.getNumElements()) + const int64_t num_factors = factor_values.getNumElements(); + if (static_cast(scales.size()) != num_factors && num_factors != 1) return {}; SmallVector new_scales; new_scales.reserve(scales.size()); - auto scales_iter = scales.begin(); - for (const auto& f : factor_values) { - new_scales.push_back(*scales_iter * - std::fabs(FloatAttr::getValueAsDouble(f))); - ++scales_iter; + if (num_factors == 1) { + const double factor_val = + std::fabs(FloatAttr::getValueAsDouble(*factor_values.begin())); + for (const double scale : scales) { + new_scales.push_back(scale * factor_val); + } + } else { + auto scales_iter = scales.begin(); + for (const auto& f : factor_values) { + new_scales.push_back(*scales_iter * + std::fabs(FloatAttr::getValueAsDouble(f))); + ++scales_iter; + } } // We are assuming symmetric quantization. auto new_ele_type = quant::UniformQuantizedPerAxisType::get( @@ -356,8 +365,21 @@ TypeAttr RescaleQuantizedType(const Type input, const Attribute factor) { quant::QuantizedType::castToExpressedType(input))) { return TypeAttr::get(new_type); } + } else if (const auto qtype = + dyn_cast(element_type)) { + if (factor_values.getNumElements() != 1) return {}; + const double factor_val = + std::fabs(FloatAttr::getValueAsDouble(*factor_values.begin())); + const double new_scale = qtype.getScale() * factor_val; + auto new_ele_type = quant::UniformQuantizedType::get( + qtype.getFlags(), qtype.getStorageType(), qtype.getExpressedType(), + new_scale, qtype.getZeroPoint(), qtype.getStorageTypeMin(), + qtype.getStorageTypeMax()); + if (const auto new_type = new_ele_type.castFromExpressedType( + quant::QuantizedType::castToExpressedType(input))) { + return TypeAttr::get(new_type); + } } - // Currently, we only support per-axis quantized type. return {}; } @@ -511,8 +533,8 @@ Type GetUniformQuantizedPerAxisTypeForWeight( const int dim_size = shape[quant_dim]; const int slice_size = - std::accumulate(std::next(shape.begin(), quant_dim + 1), shape.end(), 1, - std::multiplies()); + std::accumulate(std::next(shape.begin(), quant_dim + 1), shape.end(), + int64_t{1}, std::multiplies()); SmallVector mins(dim_size, std::numeric_limits::max()); SmallVector maxs(dim_size, std::numeric_limits::min()); const auto fp = dyn_cast(attr); diff --git a/tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.h b/tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.h index e039138f7c8a90..52f234d1a7ab81 100644 --- a/tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.h +++ b/tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.h @@ -43,6 +43,7 @@ limitations under the License. #include "mlir/IR/BuiltinAttributes.h" // from @llvm-project #include "mlir/IR/BuiltinTypeInterfaces.h" // from @llvm-project #include "mlir/IR/BuiltinTypes.h" // from @llvm-project +#include "mlir/IR/DialectResourceBlobManager.h" // from @llvm-project // IWYU pragma: keep #include "mlir/IR/IRMapping.h" // from @llvm-project #include "mlir/IR/Location.h" // from @llvm-project #include "mlir/IR/MLIRContext.h" // from @llvm-project @@ -51,6 +52,7 @@ limitations under the License. #include "mlir/IR/Operation.h" // from @llvm-project #include "mlir/IR/OperationSupport.h" // from @llvm-project #include "mlir/IR/PatternMatch.h" // from @llvm-project +#include "mlir/IR/TypeUtilities.h" // from @llvm-project #include "mlir/IR/Types.h" // from @llvm-project #include "mlir/IR/Value.h" // from @llvm-project #include "mlir/Support/LLVM.h" // from @llvm-project @@ -662,10 +664,26 @@ class QuantizationPattern : public RewritePattern { for (int i = 0, e = quantized_op->getNumOperands(); i < e; ++i) { auto def = quantized_op->getOperand(i).getDefiningOp(); if (auto q = llvm::dyn_cast_or_null(def)) { - DenseFPElementsAttr attr; + ElementsAttr attr; if (!matchPattern(q.getOperand(), m_Constant(&attr))) { continue; } + if (auto resourceAttr = + mlir::dyn_cast(attr)) { + if (AsmResourceBlob* blob = + resourceAttr.getRawHandle().getBlob()) { + ArrayRef ptr = blob->getData(); + if (DenseElementsAttr::isValidRawBuffer(resourceAttr.getType(), + ptr)) { + attr = DenseElementsAttr::getFromRawBuffer( + resourceAttr.getType(), ptr); + } + } + } + if (!mlir::isa( + mlir::getElementTypeOrSelf(attr.getType()))) { + continue; + } auto cst = arith::ConstantOp::create(rewriter, quantized_op->getLoc(), attr); quantizing_op->setOperand(i, cst.getResult()); diff --git a/tensorflow/compiler/mlir/lite/quantization/ir/QuantizeUtils.cc b/tensorflow/compiler/mlir/lite/quantization/ir/QuantizeUtils.cc index af0d21594ae957..7bc5887455258f 100644 --- a/tensorflow/compiler/mlir/lite/quantization/ir/QuantizeUtils.cc +++ b/tensorflow/compiler/mlir/lite/quantization/ir/QuantizeUtils.cc @@ -15,15 +15,31 @@ limitations under the License. #include "tensorflow/compiler/mlir/lite/quantization/ir/QuantizeUtils.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "llvm/ADT/STLExtras.h" #include "mlir/Dialect/Quant/IR/QuantTypes.h" // from @llvm-project +#include "mlir/IR/AsmState.h" // from @llvm-project #include "mlir/IR/Attributes.h" // from @llvm-project #include "mlir/IR/BuiltinAttributes.h" // from @llvm-project +#include "mlir/IR/BuiltinDialect.h" // from @llvm-project #include "mlir/IR/BuiltinTypeInterfaces.h" // from @llvm-project +#include "mlir/IR/DialectResourceBlobManager.h" // from @llvm-project #include "mlir/Support/LLVM.h" // from @llvm-project #include "tensorflow/compiler/mlir/quantization/common/ir/UniformSupport.h" -using namespace mlir; -using namespace mlir::quantfork; +namespace mlir { +namespace quantfork { /// Converts a possible primitive, real expressed value attribute to a /// corresponding storage attribute (typically FloatAttr -> IntegerAttr). @@ -32,10 +48,10 @@ using namespace mlir::quantfork; /// Returns a converter Attribute or nullptr if conversion is not possible. static Attribute convertPrimitiveValueAttr( Attribute origRealValue, quant::QuantizedType quantizedElementType, - const mlir::quant::ir::UniformQuantizedValueConverter &converter, - Type &outConvertedType) { + const mlir::quant::ir::UniformQuantizedValueConverter& converter, + Type& outConvertedType) { if (mlir::isa(origRealValue)) { - FloatAttr floatAttr = mlir::cast(origRealValue); + const FloatAttr floatAttr = mlir::cast(origRealValue); outConvertedType = quantizedElementType.getStorageType(); return IntegerAttr::get(quantizedElementType.getStorageType(), converter.quantizeFloatToInt(floatAttr.getValue())); @@ -50,10 +66,10 @@ static Attribute convertPrimitiveValueAttr( static DenseElementsAttr convertDenseFPElementsAttr( DenseFPElementsAttr realFPElementsAttr, quant::QuantizedType quantizedElementType, - const mlir::quant::ir::UniformQuantizedValueConverter &converter) { + const mlir::quant::ir::UniformQuantizedValueConverter& converter) { return realFPElementsAttr.mapValues( quantizedElementType.getStorageType(), - [&converter](const APFloat &realVal) { + [&converter](const APFloat& realVal) { return converter.quantizeFloatToInt(realVal); }); } @@ -64,12 +80,12 @@ static DenseElementsAttr convertDenseFPElementsAttr( static SparseElementsAttr convertSparseElementsAttr( SparseElementsAttr realSparseAttr, quant::QuantizedType quantizedElementType, - const mlir::quant::ir::UniformQuantizedValueConverter &converter) { + const mlir::quant::ir::UniformQuantizedValueConverter& converter) { DenseElementsAttr realDenseAttr = realSparseAttr.getValues(); if (!mlir::isa(realDenseAttr)) { return nullptr; } - DenseElementsAttr quantDenseAttr = + const DenseElementsAttr quantDenseAttr = convertDenseFPElementsAttr(mlir::cast(realDenseAttr), quantizedElementType, converter); if (!quantDenseAttr) { @@ -78,7 +94,7 @@ static SparseElementsAttr convertSparseElementsAttr( // Cast from an expressed-type-based type to storage-type-based type, // preserving the sparse shape (i.e. tensor<4xf32> -> tensor<4xi8>). - ShapedType newSparseType = mlir::dyn_cast_or_null( + const ShapedType newSparseType = mlir::dyn_cast_or_null( quantizedElementType.castExpressedToStorageType( realSparseAttr.getType())); if (!newSparseType) { @@ -88,17 +104,294 @@ static SparseElementsAttr convertSparseElementsAttr( quantDenseAttr); } +static Attribute quantizeResourceAttrPerAxisLegacy( + DenseResourceElementsAttr resourceAttr, + quant::UniformQuantizedPerAxisType quantizedElementType, + Type& outConvertedType) { + const ShapedType type = resourceAttr.getType(); + const int32_t quantDim = quantizedElementType.getQuantizedDimension(); + const uint32_t storageBitWidth = + quantizedElementType.getStorageTypeIntegralWidth(); + const bool isSigned = quantizedElementType.isSigned(); + const ArrayRef scales = quantizedElementType.getScales(); + const ArrayRef zeroPoints = quantizedElementType.getZeroPoints(); + const double clampMin = + static_cast(quantizedElementType.getStorageTypeMin()); + const double clampMax = + static_cast(quantizedElementType.getStorageTypeMax()); + + const std::string newKey = + (llvm::Twine(resourceAttr.getRawHandle().getKey()) + "_quant_axis_" + + llvm::Twine(quantDim) + "_w_" + llvm::Twine(storageBitWidth)) + .str(); + + const Type storageElemType = + IntegerType::get(resourceAttr.getContext(), storageBitWidth, + isSigned ? IntegerType::Signed : IntegerType::Signless); + const auto resType = RankedTensorType::get(type.getShape(), storageElemType); + + auto& manager = DenseResourceElementsHandle::getManagerInterface( + resourceAttr.getContext()); + if (const auto* entry = manager.getBlobManager().lookup(newKey)) { + if (entry->getBlob()) { + auto* dialect = + resourceAttr.getContext()->getLoadedDialect(); + const DenseResourceElementsHandle handle( + const_cast(entry), dialect); + outConvertedType = resType; + return DenseResourceElementsAttr::get(resType, handle); + } + } + + const AsmResourceBlob* blob = resourceAttr.getRawHandle().getBlob(); + if (!blob && resourceAttr.getRawHandle().getResource()) { + blob = resourceAttr.getRawHandle().getResource()->getBlob(); + } + if (!blob || blob->getData().empty()) return nullptr; + + const size_t numElements = type.getNumElements(); + const size_t elemByteSize = std::max(1, storageBitWidth / 8); + const size_t outByteSize = (storageBitWidth == 4) ? (numElements + 1) / 2 + : (storageBitWidth == 2) + ? (numElements + 3) / 4 + : numElements * elemByteSize; + + auto rawOutputBlob = mlir::HeapAsmResourceBlob::allocate( + outByteSize, /*align=*/64, /*dataIsMutable=*/true); + + const std::size_t dimSize = type.getDimSize(quantDim); + if (dimSize != scales.size()) { + return nullptr; + } + SmallVector converters; + converters.reserve(dimSize); + for (int i = 0, e = dimSize; i != e; ++i) { + converters.emplace_back(scales[i], zeroPoints[i], APFloat(clampMin), + APFloat(clampMax), storageBitWidth, isSigned); + } + + const auto shape = type.getShape(); + const int64_t chunkSize = + std::accumulate(std::next(shape.begin(), quantDim + 1), shape.end(), + int64_t{1}, std::multiplies()); + + const ArrayRef rawFloat( + reinterpret_cast(blob->getData().data()), numElements); + char* outData = const_cast(rawOutputBlob.getDataAs().data()); + + if (storageBitWidth == 8) { + const MutableArrayRef outInt8(reinterpret_cast(outData), + numElements); + for (size_t elemIdx = 0; elemIdx < numElements; ++elemIdx) { + const int chunkIndex = (elemIdx / chunkSize) % dimSize; + const APFloat old(rawFloat[elemIdx]); + const APInt q = converters[chunkIndex].quantizeFloatToInt(old); + outInt8[elemIdx] = static_cast(q.getSExtValue()); + } + } else if (storageBitWidth == 4) { + const MutableArrayRef outInt8(reinterpret_cast(outData), + outByteSize); + llvm::fill(outInt8, int8_t{0}); + for (size_t elemIdx = 0; elemIdx < numElements; ++elemIdx) { + const int chunkIndex = (elemIdx / chunkSize) % dimSize; + const APFloat old(rawFloat[elemIdx]); + const APInt q = converters[chunkIndex].quantizeFloatToInt(old); + const int8_t val = static_cast(q.getSExtValue()) & 0x0F; + if (elemIdx % 2 == 0) { + outInt8[elemIdx / 2] |= val; + } else { + outInt8[elemIdx / 2] |= (val << 4); + } + } + } else if (storageBitWidth == 2) { + const MutableArrayRef outInt8(reinterpret_cast(outData), + outByteSize); + llvm::fill(outInt8, int8_t{0}); + for (size_t elemIdx = 0; elemIdx < numElements; ++elemIdx) { + const int chunkIndex = (elemIdx / chunkSize) % dimSize; + const APFloat old(rawFloat[elemIdx]); + const APInt q = converters[chunkIndex].quantizeFloatToInt(old); + const int8_t val = static_cast(q.getSExtValue()) & 0x03; + const int shift = (elemIdx % 4) * 2; + outInt8[elemIdx / 4] |= (val << shift); + } + } else { + const MutableArrayRef outInt16(reinterpret_cast(outData), + numElements); + for (size_t elemIdx = 0; elemIdx < numElements; ++elemIdx) { + const int chunkIndex = (elemIdx / chunkSize) % dimSize; + const APFloat old(rawFloat[elemIdx]); + const APInt q = converters[chunkIndex].quantizeFloatToInt(old); + outInt16[elemIdx] = static_cast(q.getSExtValue()); + } + } + + outConvertedType = resType; + return DenseResourceElementsAttr::get(resType, newKey, + std::move(rawOutputBlob)); +} + +static Attribute quantizeResourceAttrPerAxisFast( + DenseResourceElementsAttr resourceAttr, + quant::UniformQuantizedPerAxisType quantizedElementType, + Type& outConvertedType) { + const ShapedType type = resourceAttr.getType(); + const int32_t quantDim = quantizedElementType.getQuantizedDimension(); + const uint32_t storageBitWidth = + quantizedElementType.getStorageTypeIntegralWidth(); + const bool isSigned = quantizedElementType.isSigned(); + const ArrayRef scales = quantizedElementType.getScales(); + const ArrayRef zeroPoints = quantizedElementType.getZeroPoints(); + const double clampMin = + static_cast(quantizedElementType.getStorageTypeMin()); + const double clampMax = + static_cast(quantizedElementType.getStorageTypeMax()); + + const std::string newKey = + (llvm::Twine(resourceAttr.getRawHandle().getKey()) + "_quant_axis_" + + llvm::Twine(quantDim) + "_w_" + llvm::Twine(storageBitWidth)) + .str(); + + const Type storageElemType = + IntegerType::get(resourceAttr.getContext(), storageBitWidth, + isSigned ? IntegerType::Signed : IntegerType::Signless); + const auto resType = RankedTensorType::get(type.getShape(), storageElemType); + + auto& manager = DenseResourceElementsHandle::getManagerInterface( + resourceAttr.getContext()); + if (const auto* entry = manager.getBlobManager().lookup(newKey)) { + if (entry->getBlob()) { + auto* dialect = + resourceAttr.getContext()->getLoadedDialect(); + const DenseResourceElementsHandle handle( + const_cast(entry), dialect); + outConvertedType = resType; + return DenseResourceElementsAttr::get(resType, handle); + } + } + + const AsmResourceBlob* blob = resourceAttr.getRawHandle().getBlob(); + if (!blob && resourceAttr.getRawHandle().getResource()) { + blob = resourceAttr.getRawHandle().getResource()->getBlob(); + } + if (!blob || blob->getData().empty()) return nullptr; + + const size_t numElements = type.getNumElements(); + const size_t elemByteSize = std::max(1, storageBitWidth / 8); + const size_t outByteSize = (storageBitWidth == 4) ? (numElements + 1) / 2 + : (storageBitWidth == 2) + ? (numElements + 3) / 4 + : numElements * elemByteSize; + + auto rawOutputBlob = mlir::HeapAsmResourceBlob::allocate( + outByteSize, /*align=*/64, /*dataIsMutable=*/true); + + const std::size_t dimSize = type.getDimSize(quantDim); + if (dimSize != scales.size()) { + return nullptr; + } + + std::vector invScales(dimSize); + std::vector zps(dimSize); + for (size_t i = 0; i < dimSize; ++i) { + invScales[i] = 1.0f / static_cast(scales[i]); + zps[i] = static_cast(zeroPoints[i]); + } + const float fMin = static_cast(clampMin); + const float fMax = static_cast(clampMax); + + const auto shape = type.getShape(); + const int64_t chunkSize = + std::accumulate(std::next(shape.begin(), quantDim + 1), shape.end(), + int64_t{1}, std::multiplies()); + + const ArrayRef rawFloat( + reinterpret_cast(blob->getData().data()), numElements); + char* outData = const_cast(rawOutputBlob.getDataAs().data()); + + if (storageBitWidth == 8) { + const MutableArrayRef outInt8(reinterpret_cast(outData), + numElements); + for (size_t elemIdx = 0; elemIdx < numElements; ++elemIdx) { + const int chunkIndex = (elemIdx / chunkSize) % dimSize; + const float val = + std::clamp(std::nearbyint(rawFloat[elemIdx] * invScales[chunkIndex]) + + zps[chunkIndex], + fMin, fMax); + outInt8[elemIdx] = static_cast(val); + } + } else if (storageBitWidth == 4) { + const MutableArrayRef outInt8(reinterpret_cast(outData), + outByteSize); + llvm::fill(outInt8, int8_t{0}); + for (size_t elemIdx = 0; elemIdx < numElements; ++elemIdx) { + const int chunkIndex = (elemIdx / chunkSize) % dimSize; + const float val = + std::clamp(std::nearbyint(rawFloat[elemIdx] * invScales[chunkIndex]) + + zps[chunkIndex], + fMin, fMax); + const int8_t nibble = static_cast(val) & 0x0F; + if (elemIdx % 2 == 0) { + outInt8[elemIdx / 2] |= nibble; + } else { + outInt8[elemIdx / 2] |= (nibble << 4); + } + } + } else if (storageBitWidth == 2) { + const MutableArrayRef outInt8(reinterpret_cast(outData), + outByteSize); + llvm::fill(outInt8, int8_t{0}); + for (size_t elemIdx = 0; elemIdx < numElements; ++elemIdx) { + const int chunkIndex = (elemIdx / chunkSize) % dimSize; + const float val = + std::clamp(std::nearbyint(rawFloat[elemIdx] * invScales[chunkIndex]) + + zps[chunkIndex], + fMin, fMax); + const int8_t twoBit = static_cast(val) & 0x03; + const int shift = (elemIdx % 4) * 2; + outInt8[elemIdx / 4] |= (twoBit << shift); + } + } else { + const MutableArrayRef outInt16(reinterpret_cast(outData), + numElements); + for (size_t elemIdx = 0; elemIdx < numElements; ++elemIdx) { + const int chunkIndex = (elemIdx / chunkSize) % dimSize; + const float val = + std::clamp(std::nearbyint(rawFloat[elemIdx] * invScales[chunkIndex]) + + zps[chunkIndex], + fMin, fMax); + outInt16[elemIdx] = static_cast(val); + } + } + + outConvertedType = resType; + return DenseResourceElementsAttr::get(resType, newKey, + std::move(rawOutputBlob)); +} + +static Attribute quantizeResourceAttrPerAxis( + DenseResourceElementsAttr resourceAttr, + quant::UniformQuantizedPerAxisType quantizedElementType, + Type& outConvertedType, bool useLegacySlowQuantize = false) { + if (useLegacySlowQuantize) { + return quantizeResourceAttrPerAxisLegacy(resourceAttr, quantizedElementType, + outConvertedType); + } + return quantizeResourceAttrPerAxisFast(resourceAttr, quantizedElementType, + outConvertedType); +} + /// Converts a real expressed Attribute to a corresponding Attribute containing /// quantized storage values assuming the given uniform quantizedElementType and /// converter. -Attribute mlir::quantfork::quantizeAttrUniform( +Attribute quantizeAttrUniform( Attribute realValue, quant::UniformQuantizedType quantizedElementType, - const mlir::quant::ir::UniformQuantizedValueConverter &converter, - Type &outConvertedType) { + const mlir::quant::ir::UniformQuantizedValueConverter& converter, + Type& outConvertedType) { // Fork to handle different variants of constants supported. if (mlir::isa(realValue)) { // Dense tensor or vector constant. - auto converted = + const auto converted = convertDenseFPElementsAttr(mlir::cast(realValue), quantizedElementType, converter); outConvertedType = converted.getType(); @@ -106,7 +399,7 @@ Attribute mlir::quantfork::quantizeAttrUniform( } if (mlir::isa(realValue)) { // Sparse tensor or vector constant. - auto converted = + const auto converted = convertSparseElementsAttr(mlir::cast(realValue), quantizedElementType, converter); outConvertedType = converted.getType(); @@ -122,22 +415,27 @@ Attribute mlir::quantfork::quantizeAttrUniform( /// quantizedElementType.getStorageType(). /// Returns nullptr if the conversion is not supported. /// On success, stores the converted type in outConvertedType. -Attribute mlir::quantfork::quantizeAttr( - Attribute realValue, quant::QuantizedType quantizedElementType, - Type &outConvertedType) { - if (auto uniformQuantized = +Attribute quantizeAttr(Attribute realValue, + quant::QuantizedType quantizedElementType, + Type& outConvertedType) { + if (const auto uniformQuantized = mlir::dyn_cast(quantizedElementType)) { - mlir::quant::ir::UniformQuantizedValueConverter converter(uniformQuantized); + const mlir::quant::ir::UniformQuantizedValueConverter converter( + uniformQuantized); return quantizeAttrUniform(realValue, uniformQuantized, converter, outConvertedType); } - if (auto uniformQuantizedPerAxis = + if (const auto uniformQuantizedPerAxis = mlir::dyn_cast( quantizedElementType)) { + if (const auto resourceAttr = + mlir::dyn_cast(realValue)) { + return quantizeResourceAttrPerAxis(resourceAttr, uniformQuantizedPerAxis, + outConvertedType); + } mlir::quant::ir::UniformQuantizedPerAxisValueConverter converter( uniformQuantizedPerAxis); - auto converted = converter.convert(realValue); - // TODO: why we need this outConvertedType? remove it? + const auto converted = converter.convert(realValue); if (converted) { outConvertedType = converted.getType(); } @@ -145,3 +443,6 @@ Attribute mlir::quantfork::quantizeAttr( } return nullptr; } + +} // namespace quantfork +} // namespace mlir diff --git a/tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.h b/tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.h index 7ec56df1e24504..ade2a0a45e8562 100644 --- a/tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.h +++ b/tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.h @@ -20,6 +20,8 @@ limitations under the License. #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project #include "mlir/IR/BuiltinOps.h" // from @llvm-project +#include "mlir/IR/MLIRContext.h" // from @llvm-project +#include "mlir/IR/PatternMatch.h" // from @llvm-project #include "mlir/Pass/Pass.h" // from @llvm-project namespace mlir::odml { diff --git a/tensorflow/compiler/mlir/lite/tests/optimize-after-quantization.mlir b/tensorflow/compiler/mlir/lite/tests/optimize-after-quantization.mlir index 97c120796d2c5d..0e7b2ec499daf3 100644 --- a/tensorflow/compiler/mlir/lite/tests/optimize-after-quantization.mlir +++ b/tensorflow/compiler/mlir/lite/tests/optimize-after-quantization.mlir @@ -12,21 +12,40 @@ // See the License for the specific language governing permissions and // limitations under the License. // ============================================================================== -// RUN: litert-opt %s -tfl-prepare-quantize -canonicalize -tfl-quantize -canonicalize -tfl-optimize -canonicalize | FileCheck %s +// RUN: litert-opt %s -tfl-optimize -tfl-propagate-qparams -tfl-bias-quantizer -tfl-fuse-qdq -canonicalize | FileCheck %s // CHECK-LABEL: fuseMulIntoPerTensorConv2dWithQDQs func.func @fuseMulIntoPerTensorConv2dWithQDQs(%arg0: tensor<256x32x32x3xf32>) -> tensor<256x8x7x3xf32> { - %cst = arith.constant dense<1.5> : tensor<3xf32> + %cst = arith.constant dense<1.5> : tensor %cst_0 = arith.constant dense<[1.0, 2.0, 3.0]> : tensor<3xf32> %w = arith.constant dense<2.0> : tensor<3x3x3x3xf32> %q = "tfl.quantize"(%w) {qtype = tensor<3x3x3x3x!quant.uniform>} : (tensor<3x3x3x3xf32>) -> tensor<3x3x3x3x!quant.uniform> %dq = "tfl.dequantize"(%q) : (tensor<3x3x3x3x!quant.uniform>) -> tensor<3x3x3x3xf32> %0 = "tfl.conv_2d"(%arg0, %dq, %cst_0) {dilation_h_factor = 2 : i32, dilation_w_factor = 3 : i32, fused_activation_function = "NONE", padding = "SAME", stride_h = 4 : i32, stride_w = 5 : i32} : (tensor<256x32x32x3xf32>, tensor<3x3x3x3xf32>, tensor<3xf32>) -> tensor<256x8x7x3xf32> - %1 = "tfl.mul"(%0, %cst) {fused_activation_function = "NONE"} : (tensor<256x8x7x3xf32>, tensor<3xf32>) -> tensor<256x8x7x3xf32> + %1 = "tfl.mul"(%0, %cst) {fused_activation_function = "NONE"} : (tensor<256x8x7x3xf32>, tensor) -> tensor<256x8x7x3xf32> func.return %1 : tensor<256x8x7x3xf32> - // CHECK: %[[weight:.*]] = arith.constant dense<3.000000e+00> : tensor<3x3x3x3xf32> - // CHECK: %[[bias:.*]] = arith.constant dense<[1.500000e+00, 3.000000e+00, 4.500000e+00]> + // CHECK: %[[bias:.*]] = arith.constant dense<[1.500000e+00, 3.000000e+00, 4.500000e+00]> : tensor<3xf32> + // CHECK: %[[qweight:.*]] = "tfl.pseudo_qconst"() <{qtype = tensor<3x3x3x3x!quant.uniform>, value = dense<21> : tensor<3x3x3x3xi8>}> + // CHECK: %[[weight:.*]] = "tfl.dequantize"(%[[qweight]]) + // CHECK: %[[conv:.*]] = "tfl.conv_2d"(%arg0, %[[weight]], %[[bias]]) + // CHECK: return %[[conv]] : tensor<256x8x7x3xf32> +} + +// CHECK-LABEL: fuseMulIntoPerAxisConv2dWithQDQsBroadcastScalar +func.func @fuseMulIntoPerAxisConv2dWithQDQsBroadcastScalar(%arg0: tensor<256x32x32x3xf32>) -> tensor<256x8x7x3xf32> { + %cst = arith.constant dense<1.5> : tensor + %cst_0 = arith.constant dense<[1.0, 2.0, 3.0]> : tensor<3xf32> + %w = arith.constant dense<2.0> : tensor<3x3x3x3xf32> + %q = "tfl.quantize"(%w) {qtype = tensor<3x3x3x3x!quant.uniform>} : (tensor<3x3x3x3xf32>) -> tensor<3x3x3x3x!quant.uniform> + %dq = "tfl.dequantize"(%q) : (tensor<3x3x3x3x!quant.uniform>) -> tensor<3x3x3x3xf32> + %0 = "tfl.conv_2d"(%arg0, %dq, %cst_0) {dilation_h_factor = 2 : i32, dilation_w_factor = 3 : i32, fused_activation_function = "NONE", padding = "SAME", stride_h = 4 : i32, stride_w = 5 : i32} : (tensor<256x32x32x3xf32>, tensor<3x3x3x3xf32>, tensor<3xf32>) -> tensor<256x8x7x3xf32> + %1 = "tfl.mul"(%0, %cst) {fused_activation_function = "NONE"} : (tensor<256x8x7x3xf32>, tensor) -> tensor<256x8x7x3xf32> + func.return %1 : tensor<256x8x7x3xf32> + + // CHECK: %[[bias:.*]] = arith.constant dense<[1.500000e+00, 3.000000e+00, 4.500000e+00]> : tensor<3xf32> + // CHECK: %[[qweight:.*]] = "tfl.pseudo_qconst"() <{qtype = tensor<3x3x3x3x!quant.uniform>, value = dense<20> : tensor<3x3x3x3xi8>}> + // CHECK: %[[weight:.*]] = "tfl.dequantize"(%[[qweight]]) // CHECK: %[[conv:.*]] = "tfl.conv_2d"(%arg0, %[[weight]], %[[bias]]) // CHECK: return %[[conv]] : tensor<256x8x7x3xf32> } diff --git a/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.cc b/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.cc index 03b7e46f2d44ea..fd6149ad9d54fb 100644 --- a/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.cc +++ b/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.cc @@ -17,9 +17,7 @@ limitations under the License. #include #include -#include -#include "llvm/Support/Casting.h" #include "mlir/Dialect/Quant/IR/QuantTypes.h" // from @llvm-project #include "mlir/IR/Builders.h" // from @llvm-project #include "mlir/IR/BuiltinAttributes.h" // from @llvm-project @@ -30,91 +28,9 @@ limitations under the License. #include "mlir/IR/Types.h" // from @llvm-project #include "mlir/Support/LLVM.h" // from @llvm-project #include "mlir/Support/LogicalResult.h" // from @llvm-project -#include "stablehlo/dialect/StablehloOps.h" // from @stablehlo namespace mlir::TFL { -LogicalResult FillCompositeParams(stablehlo::CompositeOp op, - SmallVector& scales, - SmallVector& zero_points, - int& num_bits, bool& is_signed, - bool& is_narrow_range) { - auto scale_attr = llvm::dyn_cast_or_null( - op.getCompositeAttributes().get("scale")); - if (scale_attr == nullptr) { - return failure(); - } - for (auto float_attr : scale_attr.getValues()) { - scales.push_back(float_attr.getValue().convertToDouble()); - } - - auto zero_point_attr = llvm::dyn_cast_or_null( - op.getCompositeAttributes().get("zero_point")); - if (zero_point_attr == nullptr) { - for (int i = 0; i < scales.size(); ++i) { - zero_points.push_back(0); - } - } else if (zero_point_attr.isSplat()) { - for (int i = 0; i < scales.size(); ++i) { - zero_points.push_back( - zero_point_attr.getSplatValue().getInt()); - } - } else { - for (IntegerAttr zp : zero_point_attr.getValues()) { - zero_points.push_back(zp.getInt()); - } - } - - auto dtype_attr = llvm::dyn_cast_or_null( - op.getCompositeAttributes().get("dtype")); - if (dtype_attr == nullptr) { - return failure(); - } - std::string dtype = dtype_attr.getValue().str(); - if (dtype == "i2") { - num_bits = 2; - is_signed = true; - } else if (dtype == "i4") { - num_bits = 4; - is_signed = true; - } else if (dtype == "ui4") { - num_bits = 4; - is_signed = false; - } else if (dtype == "i8") { - num_bits = 8; - is_signed = true; - } else if (dtype == "i16") { - num_bits = 16; - is_signed = true; - } else { - return failure(); - } - auto narrow_range_attr = llvm::dyn_cast_or_null( - op.getCompositeAttributes().get("narrow_range")); - if (narrow_range_attr == nullptr) { - return failure(); - } - is_narrow_range = narrow_range_attr.getValue(); - - return success(); -} - -bool IsDrqFakeQuant(stablehlo::CompositeOp op) { - if (op.getName() != "quant.fake_quant") { - return false; - } - SmallVector scales; - SmallVector zero_points; - int num_bits; - bool is_signed; - bool is_narrow_range; - if (failed(FillCompositeParams(op, scales, zero_points, num_bits, is_signed, - is_narrow_range))) { - return false; - } - return scales.empty() && zero_points.empty(); -} - LogicalResult GetStorageParams(unsigned num_bits, bool narrow_range, bool is_signed, MLIRContext* ctx, Type& storage_type, int64_t& qmin, diff --git a/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.h b/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.h index 35d2f206ad8395..a7ebd380742df7 100644 --- a/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.h +++ b/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.h @@ -18,7 +18,9 @@ limitations under the License. #include +#include "llvm/Support/Casting.h" #include "mlir/IR/Builders.h" // from @llvm-project +#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project #include "mlir/IR/Location.h" // from @llvm-project #include "mlir/IR/MLIRContext.h" // from @llvm-project #include "mlir/IR/Types.h" // from @llvm-project @@ -28,13 +30,172 @@ limitations under the License. namespace mlir::TFL { -LogicalResult FillCompositeParams(stablehlo::CompositeOp op, +template +LogicalResult FillCompositeParams(CompositeOpType op, SmallVector& scales, SmallVector& zero_points, int& num_bits, bool& is_signed, - bool& is_narrow_range); + bool& is_narrow_range) { + auto scale_attr = llvm::dyn_cast_or_null( + op.getCompositeAttributes().get("scale")); + if (scale_attr == nullptr) { + return failure(); + } + for (auto float_attr : scale_attr.template getValues()) { + scales.push_back(float_attr.getValue().convertToDouble()); + } -bool IsDrqFakeQuant(stablehlo::CompositeOp op); + auto zero_point_attr = llvm::dyn_cast_or_null( + op.getCompositeAttributes().get("zero_point")); + if (zero_point_attr == nullptr) { + for (int i = 0; i < scales.size(); ++i) { + zero_points.push_back(0); + } + } else if (zero_point_attr.isSplat()) { + for (int i = 0; i < scales.size(); ++i) { + zero_points.push_back( + zero_point_attr.template getSplatValue().getInt()); + } + } else { + for (IntegerAttr zp : zero_point_attr.template getValues()) { + zero_points.push_back(zp.getInt()); + } + } + + auto dtype_attr = llvm::dyn_cast_or_null( + op.getCompositeAttributes().get("dtype")); + if (dtype_attr == nullptr) { + return failure(); + } + auto dtype = dtype_attr.getValue(); + + if (dtype == "i2") { + num_bits = 2; + is_signed = true; + } else if (dtype == "i4") { + num_bits = 4; + is_signed = true; + } else if (dtype == "ui4") { + num_bits = 4; + is_signed = false; + } else if (dtype == "i8") { + num_bits = 8; + is_signed = true; + } else if (dtype == "i16") { + num_bits = 16; + is_signed = true; + } else { + return failure(); + } + auto narrow_range_attr = llvm::dyn_cast_or_null( + op.getCompositeAttributes().get("narrow_range")); + if (narrow_range_attr == nullptr) { + return failure(); + } + is_narrow_range = narrow_range_attr.getValue(); + + return success(); +} + +template +bool IsDrqFakeQuant(CompositeOpType op) { + if (op.getName() != "quant.fake_quant") { + return false; + } + SmallVector scales; + SmallVector zero_points; + int num_bits; + bool is_signed; + bool is_narrow_range; + if (failed(FillCompositeParams(op, scales, zero_points, num_bits, is_signed, + is_narrow_range))) { + return false; + } + return scales.empty() && zero_points.empty(); +} + +template <> +inline LogicalResult FillCompositeParams( + stablehlo::CustomCallOp op, SmallVector& scales, + SmallVector& zero_points, int& num_bits, bool& is_signed, + bool& is_narrow_range) { + auto scale_attr = + llvm::dyn_cast_or_null(op->getAttr("scale")); + if (scale_attr == nullptr) { + return failure(); + } + for (auto float_attr : scale_attr.getValues()) { + scales.push_back(float_attr.getValue().convertToDouble()); + } + + auto zero_point_attr = + llvm::dyn_cast_or_null(op->getAttr("zero_point")); + if (zero_point_attr == nullptr) { + for (int i = 0; i < scales.size(); ++i) { + zero_points.push_back(0); + } + } else if (zero_point_attr.isSplat()) { + for (int i = 0; i < scales.size(); ++i) { + zero_points.push_back( + zero_point_attr.getSplatValue().getInt()); + } + } else { + for (IntegerAttr zp : zero_point_attr.getValues()) { + zero_points.push_back(zp.getInt()); + } + } + + auto dtype_attr = llvm::dyn_cast_or_null(op->getAttr("dtype")); + if (dtype_attr == nullptr) { + return failure(); + } + auto dtype = dtype_attr.getValue(); + + if (dtype == "i2") { + num_bits = 2; + is_signed = true; + } else if (dtype == "i4") { + num_bits = 4; + is_signed = true; + } else if (dtype == "ui4") { + num_bits = 4; + is_signed = false; + } else if (dtype == "i8") { + num_bits = 8; + is_signed = true; + } else if (dtype == "i16") { + num_bits = 16; + is_signed = true; + } else { + return failure(); + } + auto narrow_range_attr = + llvm::dyn_cast_or_null(op->getAttr("narrow_range")); + if (narrow_range_attr == nullptr) { + return failure(); + } + is_narrow_range = narrow_range_attr.getValue(); + + return success(); +} + +template <> +inline bool IsDrqFakeQuant( + stablehlo::CustomCallOp op) { + if (op.getCallTargetName() != "quant.fake_quant") { + return false; + } + SmallVector scales; + SmallVector zero_points; + int num_bits; + bool is_signed; + bool is_narrow_range; + if (failed(FillCompositeParams(op, scales, zero_points, num_bits, is_signed, + is_narrow_range))) { + return false; + } + return scales.empty() && zero_points.empty(); +} LogicalResult GetStorageParams(unsigned num_bits, bool narrow_range, bool is_signed, MLIRContext* ctx, diff --git a/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_pass.cc b/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_pass.cc index 272ce6c866436f..9ecf7dca971e65 100644 --- a/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_pass.cc +++ b/tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_pass.cc @@ -20,9 +20,8 @@ limitations under the License. #include #include "llvm/Support/Casting.h" -#include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project #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/Quant.h" // from @llvm-project // IWYU pragma: keep #include "mlir/IR/Builders.h" // from @llvm-project #include "mlir/IR/BuiltinAttributes.h" // from @llvm-project #include "mlir/IR/BuiltinOps.h" // from @llvm-project @@ -41,7 +40,6 @@ limitations under the License. #include "tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.h" #include "tensorflow/compiler/mlir/lite/transforms/passes.h" // IWYU pragma: keep #include "tensorflow/compiler/mlir/lite/utils/utils.h" -#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h" #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h" namespace mlir { @@ -235,15 +233,16 @@ class RewriteDequantizeCompositeOp Value operand = composite_op.getOperand(num_operands - 1); mlir::Operation* producer_op = operand.getDefiningOp(); - // Check if the producer is an arith.constant - if (auto const_op = - llvm::dyn_cast_or_null(producer_op)) { - // We found a constant (Int4/Int8). + // Check if the producer is a constant (arith.constant or + // stablehlo.constant) + if (producer_op && producer_op->hasAttr("value")) { + // We found a constant (Int4/Int8) with ElementsAttr (DenseElementsAttr + // or DenseResourceElementsAttr). // Instead of casting or hacking the constant, we create a valid // TFL::QConstOp. This op natively maps "Integer Data" -> "Quantized // Type". - auto value_attr = llvm::dyn_cast(const_op.getValue()); + auto value_attr = producer_op->getAttrOfType("value"); if (!value_attr) { return failure(); // Should not happen for tensor constants } @@ -251,7 +250,7 @@ class RewriteDequantizeCompositeOp // Create tfl.qconst // Arguments: Type (Result), TypeAttr (qtype), ElementsAttr (value) auto qconst_op = rewriter.create( - const_op.getLoc(), + producer_op->getLoc(), qtensor_type, // The Result Type (!quant.uniform...) TypeAttr::get(qtensor_type), // The "qtype" attribute value_attr // The reuse of the i4/i8 data @@ -260,7 +259,7 @@ class RewriteDequantizeCompositeOp // Use the output of this new constant tfl_quantize_input = qconst_op.getResult(); - // Note: We leave the old arith.constant alone. + // Note: We leave the old constant alone. // If it has no other uses, the cleanup pass (DCE) will remove it // automatically. @@ -439,9 +438,277 @@ void LowerQuantAnnotationsPass::runOnOperation() { signalPassFailure(); } + class RewriteQuantizeCustomCallOp + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(stablehlo::CustomCallOp op, + PatternRewriter& rewriter) const final { + if (op.getCallTargetName() != "quant.quantize") { + return failure(); + } + + SmallVector scales; + SmallVector zero_points; + int num_bits; + bool is_signed; + bool is_narrow_range; + + if (failed(FillCompositeParams(op, scales, zero_points, num_bits, + is_signed, is_narrow_range))) { + return op.emitError( + "quantize custom call does not contain the required attributes."); + } + + ShapedType input_shaped_type = + cast(op.getOperand(0).getType()); + Type input_element_type = input_shaped_type.getElementType(); + + Type quantized_element_type; + if (scales.size() == 1) { + quantized_element_type = GetPerTensorQuantizedTensorType( + rewriter, scales[0], zero_points[0], + /*expressed_type=*/input_element_type, num_bits, op->getLoc(), + is_narrow_range, is_signed); + } else { + int32_t quantized_dimension; + if (auto quantized_dimension_attr = llvm::dyn_cast_or_null( + op->getAttr("quantization_dimension"))) { + quantized_dimension = + quantized_dimension_attr.getValue().getSExtValue(); + } else { + return op.emitError( + "quantization_dimension attribute is missing from the custom " + "call."); + } + quantized_element_type = GetPerAxisQuantizedTensorType( + rewriter, scales, zero_points, quantized_dimension, + /*expressed_type=*/input_element_type, num_bits, op->getLoc(), + is_narrow_range, is_signed); + } + + RankedTensorType output_type = RankedTensorType::get( + input_shaped_type.getShape(), quantized_element_type); + TFL::QuantizeOp tfl_quantize_op = + TFL::QuantizeOp::create(rewriter, op.getLoc(), output_type, + /*input=*/op.getOperand(0), + /*qtype=*/TypeAttr::get(output_type)); + + rewriter.replaceOp(op, tfl_quantize_op.getOutput()); + return success(); + } + }; + + class RewriteDequantizeCustomCallOp + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(stablehlo::CustomCallOp custom_call_op, + PatternRewriter& rewriter) const final { + if (custom_call_op.getCallTargetName() != "quant.dequantize") { + return failure(); + } + Type output_type = custom_call_op.getType(0); + + SmallVector scales; + SmallVector zero_points; + int num_bits; + bool is_signed; + bool is_narrow_range; + + if (failed(FillCompositeParams(custom_call_op, scales, zero_points, + num_bits, is_signed, is_narrow_range))) { + return failure(); + } + + ShapedType output_shaped_type = cast(output_type); + Type output_element_type = output_shaped_type.getElementType(); + + Type quantized_element_type; + if (scales.size() == 1) { + quantized_element_type = GetPerTensorQuantizedTensorType( + rewriter, scales[0], zero_points[0], + /*expressed_type=*/output_element_type, num_bits, + custom_call_op->getLoc(), is_narrow_range, is_signed); + } else { + int32_t quantized_dimension; + if (auto quantized_dimension_attr = llvm::dyn_cast_or_null( + custom_call_op->getAttr("quantization_dimension"))) { + quantized_dimension = + quantized_dimension_attr.getValue().getSExtValue(); + } else { + return failure(); + } + quantized_element_type = GetPerAxisQuantizedTensorType( + rewriter, scales, zero_points, quantized_dimension, + /*expressed_type=*/output_element_type, num_bits, + custom_call_op->getLoc(), is_narrow_range, is_signed); + } + + ShapedType input_shaped_type = + cast(custom_call_op.getOperand(0).getType()); + RankedTensorType qtensor_type = RankedTensorType::get( + input_shaped_type.getShape(), quantized_element_type); + + auto custom_call_operand = custom_call_op.getOperand(0); + + Value tfl_quantize_input; + if (mlir::dyn_cast(custom_call_operand)) { + // Find the function enclosing this custom call op. + func::FuncOp func_op = GetEnclosingFunction(custom_call_op); + if (func_op == nullptr) { + return failure(); + } + + // Find the operand index of the input of the custom call op. + int arg_idx = -1; + for (int i = 0; i < func_op.getNumArguments(); ++i) { + if (func_op.getBody().front().getArgument(i) == custom_call_operand) { + arg_idx = i; + break; + } + } + if (arg_idx == -1) { + return failure(); + } + + // create a new set of operand types for the function with the type of + // the operand that feeds the custom call op changed. + SmallVector new_func_input_types; + auto func_input_types = func_op.getFunctionType().getInputs(); + for (int i = 0; i < func_input_types.size(); ++i) { + if (i != arg_idx) { + new_func_input_types.push_back(func_input_types[i]); + } else { + new_func_input_types.push_back(qtensor_type); + } + } + + auto new_func_type = + mlir::FunctionType::get(func_op.getContext(), new_func_input_types, + func_op.getFunctionType().getResults()); + + rewriter.startOpModification(func_op); + // Update the function type. + func_op.setType(new_func_type); + + // Update the block argument type. + func_op.getBody().front().getArgument(arg_idx).setType(qtensor_type); + rewriter.finalizeOpModification(func_op); + + tfl_quantize_input = func_op.getBody().front().getArgument(arg_idx); + } else { + // Get the producer of the input to dequantize + int num_operands = custom_call_op.getNumOperands(); + Value operand = custom_call_op.getOperand(num_operands - 1); + mlir::Operation* producer_op = operand.getDefiningOp(); + + if (producer_op && producer_op->hasAttr("value")) { + auto value_attr = producer_op->getAttrOfType("value"); + if (!value_attr) { + return failure(); + } + + auto qconst_op = rewriter.create( + producer_op->getLoc(), qtensor_type, TypeAttr::get(qtensor_type), + value_attr); + + tfl_quantize_input = qconst_op.getResult(); + } else if (producer_op) { + rewriter.startOpModification(producer_op); + for (OpResult result : producer_op->getResults()) { + if (result == custom_call_op.getOperand(num_operands - 1)) { + result.setType(qtensor_type); + break; + } + } + rewriter.finalizeOpModification(producer_op); + + tfl_quantize_input = operand; + } else { + return failure(); + } + } + + TFL::DequantizeOp tfl_dequantize_op = TFL::DequantizeOp::create( + rewriter, custom_call_op.getLoc(), output_type, + /*input=*/tfl_quantize_input); + rewriter.replaceOp(custom_call_op, tfl_dequantize_op.getOutput()); + return success(); + } + }; + + class RewriteFakeQuantCustomCallOp + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + public: + LogicalResult matchAndRewrite(stablehlo::CustomCallOp op, + PatternRewriter& rewriter) const final { + if (op.getCallTargetName() != "quant.fake_quant" || IsDrqFakeQuant(op)) { + return failure(); + } + + SmallVector scales; + SmallVector zero_points; + int num_bits; + bool is_signed; + bool is_narrow_range; + + if (failed(FillCompositeParams(op, scales, zero_points, num_bits, + is_signed, is_narrow_range))) { + return op.emitError( + "fake quant custom call does not contain the required attributes."); + } + + ShapedType input_shaped_type = + cast(op.getOperand(0).getType()); + Type input_element_type = input_shaped_type.getElementType(); + + Type quantized_element_type; + if (scales.size() == 1) { + quantized_element_type = GetPerTensorQuantizedTensorType( + rewriter, scales[0], zero_points[0], + /*expressed_type=*/input_element_type, num_bits, op->getLoc(), + is_narrow_range, is_signed); + } else { + int32_t quantized_dimension; + if (auto quantized_dimension_attr = llvm::dyn_cast_or_null( + op->getAttr("quantization_dimension"))) { + quantized_dimension = + quantized_dimension_attr.getValue().getSExtValue(); + } else { + return op.emitError( + "quantization_dimension attribute is missing from the custom " + "call."); + } + quantized_element_type = GetPerAxisQuantizedTensorType( + rewriter, scales, zero_points, quantized_dimension, + /*expressed_type=*/input_element_type, num_bits, op->getLoc(), + is_narrow_range, is_signed); + } + + RankedTensorType output_type = RankedTensorType::get( + input_shaped_type.getShape(), quantized_element_type); + TFL::QuantizeOp tfl_quantize_op = + TFL::QuantizeOp::create(rewriter, op.getLoc(), output_type, + /*input=*/op.getOperand(0), + /*qtype=*/TypeAttr::get(output_type)); + + TFL::DequantizeOp tfl_dequantize_op = + TFL::DequantizeOp::create(rewriter, op.getLoc(), input_shaped_type, + tfl_quantize_op.getOutput()); + + rewriter.replaceOp(op, tfl_dequantize_op.getOutput()); + return success(); + } + }; + RewritePatternSet patterns(&ctx); patterns.add(&ctx); + RewriteFakeQuantCompositeOp, RewriteQuantizeCustomCallOp, + RewriteDequantizeCustomCallOp, RewriteFakeQuantCustomCallOp>( + &ctx); if (failed( applyPatternsGreedily(module, std::move(patterns), greedy_config))) { diff --git a/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc b/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc index 668493eca931e7..a3f758306afacc 100644 --- a/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc +++ b/tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc @@ -91,17 +91,21 @@ struct ConvertBatchMatMulOp2FullyConnectedOp_Rank2ConstantRhs rhs = reshape.getInput(); } - DenseElementsAttr dense_constant; - if (matchPattern(rhs, m_Constant(&dense_constant))) { - constant = dense_constant; + ElementsAttr elements_constant; + if (matchPattern(rhs, m_Constant(&elements_constant))) { + constant = elements_constant; } else if (auto dq = rhs.getDefiningOp()) { Value q_input = dq.getInput(); - if (auto q = q_input.getDefiningOp()) { - if (matchPattern(q.getInput(), m_Constant(&dense_constant))) { - constant = dense_constant; + if (matchPattern(q_input, m_Constant(&elements_constant))) { + constant = elements_constant; + } else if (auto q = q_input.getDefiningOp()) { + if (matchPattern(q.getInput(), m_Constant(&elements_constant))) { + constant = elements_constant; } } else if (auto pseudo_q = q_input.getDefiningOp()) { constant = pseudo_q.getValue(); + } else if (auto const_op = q_input.getDefiningOp()) { + constant = const_op.getValue(); } } diff --git a/tensorflow/compiler/mlir/lite/transforms/passes.h b/tensorflow/compiler/mlir/lite/transforms/passes.h index 84b7476d569fed..c9e71b5b75d6c5 100644 --- a/tensorflow/compiler/mlir/lite/transforms/passes.h +++ b/tensorflow/compiler/mlir/lite/transforms/passes.h @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include #include "absl/container/flat_hash_set.h" #include "mlir/Pass/Pass.h" // from @llvm-project diff --git a/tensorflow/compiler/mlir/lite/transforms/post_quantize.cc b/tensorflow/compiler/mlir/lite/transforms/post_quantize.cc index 853d6801cc8a75..3859bb49924cea 100644 --- a/tensorflow/compiler/mlir/lite/transforms/post_quantize.cc +++ b/tensorflow/compiler/mlir/lite/transforms/post_quantize.cc @@ -15,30 +15,45 @@ limitations under the License. // This transformation pass applies some clean up steps after quantization. +#include +#include +#include #include +#include +#include #include #include #include +#include #include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/SmallVector.h" #include "llvm/Support/Casting.h" #include "mlir/Dialect/Arith/IR/Arith.h" // from @llvm-project #include "mlir/Dialect/Quant/IR/QuantTypes.h" // from @llvm-project +#include "mlir/IR/AsmState.h" // from @llvm-project +#include "mlir/IR/Builders.h" // from @llvm-project +#include "mlir/IR/BuiltinAttributeInterfaces.h" // from @llvm-project #include "mlir/IR/BuiltinAttributes.h" // from @llvm-project #include "mlir/IR/BuiltinTypeInterfaces.h" // from @llvm-project #include "mlir/IR/BuiltinTypes.h" // from @llvm-project +#include "mlir/IR/DialectResourceBlobManager.h" // from @llvm-project #include "mlir/IR/Location.h" // from @llvm-project #include "mlir/IR/MLIRContext.h" // from @llvm-project +#include "mlir/IR/Matchers.h" // from @llvm-project +#include "mlir/IR/OpDefinition.h" // from @llvm-project #include "mlir/IR/PatternMatch.h" // from @llvm-project #include "mlir/IR/TypeUtilities.h" // from @llvm-project #include "mlir/Pass/Pass.h" // from @llvm-project #include "mlir/Support/LLVM.h" // from @llvm-project #include "mlir/Support/LogicalResult.h" // from @llvm-project +#include "mlir/Support/TypeID.h" // from @llvm-project #include "mlir/Transforms/GreedyPatternRewriteDriver.h" // from @llvm-project #include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h" #include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_config.h" #include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.h" #include "tensorflow/compiler/mlir/lite/transforms/passes.h" +#include "tensorflow/compiler/mlir/lite/utils/utils.h" #include "tensorflow/compiler/mlir/tensorflow/utils/dynamic_shape_utils.h" //===----------------------------------------------------------------------===// @@ -360,6 +375,27 @@ struct RemoveVolatileOps : public OpRewritePattern { } }; +static bool MatchElementsAttr(Value val, ElementsAttr& attr) { + if (matchPattern(val, m_Constant(&attr))) { + return true; + } + if (Operation* op = val.getDefiningOp()) { + if (auto const_op = llvm::dyn_cast(op)) { + attr = const_op.getValue(); + return true; + } + if (auto qconst_op = llvm::dyn_cast(op)) { + attr = qconst_op.getValue(); + return true; + } + if (auto arith_const = llvm::dyn_cast(op)) { + attr = mlir::dyn_cast(arith_const.getValue()); + return attr != nullptr; + } + } + return false; +} + // Fold the constant quantized Transpose ops. struct FoldTransposeOp : public OpRewritePattern { explicit FoldTransposeOp(MLIRContext* context) @@ -394,16 +430,49 @@ struct FoldTransposeOp : public OpRewritePattern { } } + void ComputePermutationRaw(ArrayRef perm, + ArrayRef output_shape, + const char* raw_input, int element_byte_size, + int output_axis, char*& raw_output, + SmallVectorImpl& current_input_index, + ArrayRef input_shape) const { + const int num_dimensions = output_shape.size(); + assert(output_axis < num_dimensions); + const int input_axis = perm[output_axis]; + for (int i = 0; i < output_shape[output_axis]; ++i) { + current_input_index[input_axis] = i; + const bool is_last_axis = output_axis == num_dimensions - 1; + if (is_last_axis) { + uint64_t input_flat_index = 0; + uint64_t stride = 1; + for (int d = num_dimensions - 1; d >= 0; --d) { + input_flat_index += current_input_index[d] * stride; + stride *= input_shape[d]; + } + memcpy(raw_output, raw_input + input_flat_index * element_byte_size, + element_byte_size); + raw_output += element_byte_size; + } else { + ComputePermutationRaw(perm, output_shape, raw_input, element_byte_size, + output_axis + 1, raw_output, current_input_index, + input_shape); + } + } + } + LogicalResult matchAndRewrite(TransposeOp op, PatternRewriter& rewriter) const override { Operation* def_op = op.getInput().getDefiningOp(); auto qconst_op = llvm::dyn_cast_or_null(def_op); if (qconst_op == nullptr) return failure(); - DenseIntElementsAttr perm_tensor; - if (!matchPattern(op.getPerm(), m_Constant(&perm_tensor))) return failure(); + ElementsAttr perm_attr; + if (!MatchElementsAttr(op.getPerm(), perm_attr)) return failure(); + auto int_perm_attr = mlir::dyn_cast(perm_attr); + if (!int_perm_attr) return failure(); - auto output_element_type = getElementTypeOrSelf(op.getOutput().getType()); + auto result_type = mlir::cast(op.getOutput().getType()); + auto output_element_type = result_type.getElementType(); if (!mlir::isa(output_element_type) && !mlir::isa(output_element_type)) { return failure(); @@ -411,48 +480,164 @@ struct FoldTransposeOp : public OpRewritePattern { ElementsAttr input_tensor = qconst_op.getValue(); - assert(perm_tensor.getType().getRank() == 1); const int num_dimensions = input_tensor.getShapedType().getRank(); - assert(perm_tensor.getType().getNumElements() == num_dimensions); - ArrayRef input_shape = input_tensor.getShapedType().getShape(); - auto output_type = mlir::cast(op.getOutput().getType()); SmallVector perm; + for (const APInt& it : int_perm_attr.getValues()) { + perm.push_back(it.getSExtValue()); + } + if (perm.size() != num_dimensions) return failure(); + SmallVector output_shape; for (int i = 0; i < num_dimensions; ++i) { - perm.push_back(perm_tensor.getValues()[i].getInt()); output_shape.push_back(input_shape[perm[i]]); + assert(!result_type.hasStaticShape() || + result_type.getShape()[i] == output_shape[i]); + } + + if (auto dense_input = mlir::dyn_cast(input_tensor)) { + std::vector new_values; + new_values.reserve(input_tensor.getShapedType().getNumElements()); + std::vector input_indices(num_dimensions); + ComputePermutation(dense_input, perm, output_shape, num_dimensions, + /*output_axis=*/0, &input_indices, &new_values); + RankedTensorType values_type; + if (mlir::isa(output_element_type)) { + values_type = RankedTensorType::get( + output_shape, + mlir::cast(output_element_type) + .getStorageType()); + } else { + values_type = RankedTensorType::get( + output_shape, + mlir::cast(output_element_type) + .getStorageType()); + } - // Check that the derived output shape matches the static shape. - assert(!output_type.hasStaticShape() || - output_type.getShape()[i] == output_shape[i]); + rewriter.replaceOpWithNewOp( + op, TypeAttr::get(result_type), + DenseIntElementsAttr::get(values_type, new_values)); + return success(); } - std::vector new_values; - new_values.reserve(input_tensor.getShapedType().getNumElements()); - std::vector input_indices(num_dimensions); - ComputePermutation(input_tensor, perm, output_shape, num_dimensions, - /*output_axis=*/0, &input_indices, &new_values); - auto result_type = - RankedTensorType::get(output_shape, output_type.getElementType()); - RankedTensorType values_type; - if (mlir::isa(output_element_type)) { - values_type = RankedTensorType::get( - output_shape, - mlir::cast(output_type.getElementType()) - .getStorageType()); - } else { - values_type = RankedTensorType::get( - output_shape, mlir::cast( - output_type.getElementType()) - .getStorageType()); + if (auto dense_res = + mlir::dyn_cast(input_tensor)) { + AsmResourceBlob* blob = dense_res.getRawHandle().getBlob(); + if (!blob && dense_res.getRawHandle().getResource()) { + blob = dense_res.getRawHandle().getResource()->getBlob(); + } + if (!blob || blob->getData().empty()) return failure(); + + uint32_t storage_bit_width = 8; + if (auto u = mlir::dyn_cast( + output_element_type)) { + storage_bit_width = u.getStorageTypeIntegralWidth(); + } else if (auto p = mlir::dyn_cast( + output_element_type)) { + storage_bit_width = p.getStorageTypeIntegralWidth(); + } + + size_t num_elements = result_type.getNumElements(); + if (storage_bit_width == 4) { + size_t out_byte_size = (num_elements + 1) / 2; + auto raw_output_blob = mlir::HeapAsmResourceBlob::allocate( + out_byte_size, /*align=*/64, /*dataIsMutable=*/true); + char* raw_output = + const_cast(raw_output_blob.getDataAs().data()); + std::memset(raw_output, 0, out_byte_size); + + const uint8_t* raw_input_u8 = + reinterpret_cast(blob->getData().data()); + std::vector unpacked_input(num_elements); + for (size_t i = 0; i < num_elements; ++i) { + uint8_t byte = raw_input_u8[i / 2]; + uint8_t nibble = (i % 2 == 0) ? (byte & 0x0F) : ((byte >> 4) & 0x0F); + unpacked_input[i] = static_cast(nibble); + } + + std::vector unpacked_output(num_elements); + char* unpacked_output_ptr = unpacked_output.data(); + SmallVector current_input_index(num_dimensions, 0); + ComputePermutationRaw(perm, output_shape, unpacked_input.data(), + /*element_byte_size=*/1, + /*output_axis=*/0, unpacked_output_ptr, + current_input_index, input_shape); + + for (size_t i = 0; i < num_elements; ++i) { + uint8_t nibble = static_cast(unpacked_output[i]) & 0x0F; + if (i % 2 == 0) { + raw_output[i / 2] |= nibble; + } else { + raw_output[i / 2] |= (nibble << 4); + } + } + + DenseResourceElementsAttr new_res_attr = DenseResourceElementsAttr::get( + result_type, dense_res.getRawHandle().getKey(), + std::move(raw_output_blob)); + rewriter.replaceOpWithNewOp(op, TypeAttr::get(result_type), + new_res_attr); + return success(); + } else if (storage_bit_width == 2) { + size_t out_byte_size = (num_elements + 3) / 4; + auto raw_output_blob = mlir::HeapAsmResourceBlob::allocate( + out_byte_size, /*align=*/64, /*dataIsMutable=*/true); + char* raw_output = + const_cast(raw_output_blob.getDataAs().data()); + std::memset(raw_output, 0, out_byte_size); + + const uint8_t* raw_input_u8 = + reinterpret_cast(blob->getData().data()); + std::vector unpacked_input(num_elements); + for (size_t i = 0; i < num_elements; ++i) { + uint8_t byte = raw_input_u8[i / 4]; + uint8_t val = (byte >> ((i % 4) * 2)) & 0x03; + unpacked_input[i] = static_cast(val); + } + + std::vector unpacked_output(num_elements); + char* unpacked_output_ptr = unpacked_output.data(); + SmallVector current_input_index(num_dimensions, 0); + ComputePermutationRaw(perm, output_shape, unpacked_input.data(), + /*element_byte_size=*/1, + /*output_axis=*/0, unpacked_output_ptr, + current_input_index, input_shape); + + for (size_t i = 0; i < num_elements; ++i) { + uint8_t val = static_cast(unpacked_output[i]) & 0x03; + raw_output[i / 4] |= (val << ((i % 4) * 2)); + } + + DenseResourceElementsAttr new_res_attr = DenseResourceElementsAttr::get( + result_type, dense_res.getRawHandle().getKey(), + std::move(raw_output_blob)); + rewriter.replaceOpWithNewOp(op, TypeAttr::get(result_type), + new_res_attr); + return success(); + } + + const int element_byte_size = std::max(1, storage_bit_width / 8); + auto raw_output_blob = mlir::HeapAsmResourceBlob::allocate( + blob->getData().size(), /*align=*/64, /*dataIsMutable=*/true); + char* raw_output = + const_cast(raw_output_blob.getDataAs().data()); + const char* raw_input = blob->getData().data(); + + SmallVector current_input_index(num_dimensions, 0); + ComputePermutationRaw(perm, output_shape, raw_input, element_byte_size, + /*output_axis=*/0, raw_output, current_input_index, + input_shape); + + DenseResourceElementsAttr new_res_attr = DenseResourceElementsAttr::get( + result_type, dense_res.getRawHandle().getKey(), + std::move(raw_output_blob)); + rewriter.replaceOpWithNewOp(op, TypeAttr::get(result_type), + new_res_attr); + return success(); } - rewriter.replaceOpWithNewOp( - op, TypeAttr::get(result_type), - DenseIntElementsAttr::get(values_type, new_values)); - return success(); + return failure(); } }; @@ -471,10 +656,6 @@ struct FoldReshapeOp : public OpRewritePattern { return rewriter.notifyMatchFailure(op, "input is not a QConstOp."); } - auto dense_elements = - mlir::dyn_cast_or_null(qconst_op.getValue()); - if (dense_elements == nullptr) return failure(); - auto output_element_type = getElementTypeOrSelf(op.getType()); if (!mlir::isa(output_element_type)) { return rewriter.notifyMatchFailure(op, "output type is not quantized."); @@ -488,9 +669,10 @@ struct FoldReshapeOp : public OpRewritePattern { // If the result type isn't static, tries to derive the result type from // the #2 operand. if (!result_type.hasStaticShape()) { - DenseIntElementsAttr shape_elements; - if (!matchPattern(op.getShape(), m_Constant(&shape_elements))) - return failure(); + ElementsAttr shape_attr; + if (!MatchElementsAttr(op.getShape(), shape_attr)) return failure(); + auto shape_elements = mlir::dyn_cast(shape_attr); + if (!shape_elements) return failure(); SmallVector shape_data; for (const APInt& it : shape_elements.getValues()) { @@ -499,24 +681,55 @@ struct FoldReshapeOp : public OpRewritePattern { result_type = RankedTensorType::get(shape_data, input_type.getElementType()); } + RankedTensorType values_type; - if (mlir::isa(output_element_type)) { + if (auto uniform_qtype = + mlir::dyn_cast(output_element_type)) { + values_type = RankedTensorType::get(result_type.getShape(), + uniform_qtype.getStorageType()); + } else { values_type = RankedTensorType::get( result_type.getShape(), - mlir::cast(result_type.getElementType()) + mlir::cast(output_element_type) .getStorageType()); - } else { - values_type = - RankedTensorType::get(result_type.getShape(), - mlir::cast( - result_type.getElementType()) - .getStorageType()); } - DenseElementsAttr reshaped_elements = dense_elements.reshape(values_type); - rewriter.replaceOpWithNewOp(op, TypeAttr::get(result_type), - reshaped_elements); - return success(); + ElementsAttr value_attr = qconst_op.getValue(); + if (auto dense_elements = mlir::dyn_cast(value_attr)) { + DenseElementsAttr reshaped_elements = dense_elements.reshape(values_type); + rewriter.replaceOpWithNewOp(op, TypeAttr::get(result_type), + reshaped_elements); + return success(); + } + + if (auto dense_resource_elements = + mlir::dyn_cast(value_attr)) { + AsmResourceBlob* blob = dense_resource_elements.getRawHandle().getBlob(); + if (!blob && dense_resource_elements.getRawHandle().getResource()) { + blob = dense_resource_elements.getRawHandle().getResource()->getBlob(); + } + if (!blob || blob->getData().empty()) return failure(); + + DenseResourceElementsAttr new_res_attr; + if (qconst_op.getOutput().hasOneUse()) { + new_res_attr = DenseResourceElementsAttr::get( + result_type, dense_resource_elements.getRawHandle().getKey(), + std::move(*blob)); + } else { + auto new_blob = mlir::HeapAsmResourceBlob::allocate( + blob->getData().size(), /*align=*/64, true); + memcpy(const_cast(new_blob.getData().data()), + blob->getData().data(), blob->getData().size()); + new_res_attr = DenseResourceElementsAttr::get( + result_type, dense_resource_elements.getRawHandle().getKey(), + std::move(new_blob)); + } + rewriter.replaceOpWithNewOp(op, TypeAttr::get(result_type), + new_res_attr); + return success(); + } + + return failure(); } }; diff --git a/tensorflow/compiler/mlir/lite/transforms/quantization/fuse_qdq_pass.cc b/tensorflow/compiler/mlir/lite/transforms/quantization/fuse_qdq_pass.cc index 79f832ad31ab50..a293d9ff9f6a7d 100644 --- a/tensorflow/compiler/mlir/lite/transforms/quantization/fuse_qdq_pass.cc +++ b/tensorflow/compiler/mlir/lite/transforms/quantization/fuse_qdq_pass.cc @@ -78,6 +78,14 @@ LogicalResult IsDrqTensor(mlir::Value value, mlir::Value& fq_input) { return success(); } } + if (auto custom_call_op = llvm::dyn_cast_or_null( + value.getDefiningOp())) { + if (IsDrqFakeQuant(custom_call_op)) { + int num_operands = custom_call_op.getNumOperands(); + fq_input = custom_call_op.getOperand(num_operands - 1); + return success(); + } + } return failure(); } @@ -368,6 +376,21 @@ class RemoveUnusedFQ : public OpRewritePattern { } }; +class RemoveUnusedCustomCallFQ + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(stablehlo::CustomCallOp op, + PatternRewriter& rewriter) const final { + if (IsDrqFakeQuant(op) && op->getUses().empty()) { + rewriter.eraseOp(op); + return success(); + } + return rewriter.notifyMatchFailure( + op, "is not a drq fake quant custom call with no uses."); + } +}; + // Pushes a drq fake quant op forward through a pad op. // This is to allow DRQ FQ to be fused into the DRQ op. // drq_fake_quant(input) -> pad -> output @@ -404,19 +427,45 @@ class PushForwardDrqFQ : public OpRewritePattern { TFL::PadOp::create(rewriter, pad_op.getLoc(), pad_op.getType(), float_input, pad_op.getPadding()); - // Create a new drq fake quant op. - // Operands are the same, except for the last one. - SmallVector new_drq_operands; - for (mlir::Value operand : drq_fq_op.getOperands().drop_back()) { - new_drq_operands.push_back(operand); - } - new_drq_operands.push_back(new_pad_op.getResult()); + Operation* new_drq_fq_op = rewriter.clone(*drq_fq_op.getOperation()); + new_drq_fq_op->setOperand(new_drq_fq_op->getNumOperands() - 1, + new_pad_op.getResult()); + + rewriter.replaceOp(pad_op, new_drq_fq_op->getResult(0)); + return success(); + } +}; + +// Fixes keep_num_dims option of FC if output dims is different from input dims +// though keep_num_dims is true. It happens when FC's input has changed after +// quantization, e.g. by IsDrqTensor(). +// Sets keep_num_dims to false if that's the case. Otherwise, it's not +// compatible with GPU. See CheckGpuDelegateCompatibility() in +// third_party/tensorflow/lite/tools/versioning/gpu_compatibility.cc. +// Note that if FC is followed by Reshape, the keep_num_dims will be set to true +// with a correct shape later by EnableFullyConnectedKeepNumDimsBeforeReshape() +// in optimize pass. +struct FixFullyConnectedKeepNumDims + : public OpRewritePattern { + explicit FixFullyConnectedKeepNumDims(MLIRContext* context) + : OpRewritePattern(context, /*benefit=*/0) {} + + LogicalResult matchAndRewrite(TFL::FullyConnectedOp fc, + PatternRewriter& rewriter) const override { + if (!fc.getKeepNumDims()) return failure(); - auto new_drq_fq_op = stablehlo::CompositeOp::create( - rewriter, drq_fq_op.getLoc(), pad_op.getType(), new_drq_operands, - drq_fq_op->getAttrs()); + auto input_ty = + mlir::dyn_cast_or_null(fc.getInput().getType()); + auto fc_ty = mlir::dyn_cast_or_null(fc.getType(0)); + if (!input_ty || !fc_ty) return failure(); + + auto input_shape = input_ty.getShape(); + auto fc_shape = fc_ty.getShape(); + if (input_shape.size() == fc_shape.size()) { + return failure(); + } - rewriter.replaceOp(pad_op, new_drq_fq_op.getResult(0)); + rewriter.modifyOpInPlace(fc, [&]() { fc.setKeepNumDims(false); }); return success(); } }; @@ -484,7 +533,7 @@ class QuantizeConstPattern : public OpRewritePattern { : OpRewritePattern(context) {} LogicalResult matchAndRewrite(QuantizeOp op, PatternRewriter& rewriter) const override { - DenseFPElementsAttr attr; + ElementsAttr attr; if (matchPattern(op.getInput(), m_Constant(&attr))) { auto qtype = op.getQtypeAttr(); Attribute quantized_attr = mlir::TFL::Quantize(attr, qtype.getValue()); @@ -524,8 +573,10 @@ void FuseQDQPass::runOnOperation() { mlir::ModuleOp module = getOperation(); RewritePatternSet patterns(ctx); - patterns.add(ctx); + patterns + .add(ctx); // Configure the greedy pattern rewrite driver. GreedyRewriteConfig greedy_config; diff --git a/tensorflow/compiler/mlir/lite/transforms/reduce_type_precision.cc b/tensorflow/compiler/mlir/lite/transforms/reduce_type_precision.cc index 865ae58a04bfde..aa51e8f92b9e1f 100644 --- a/tensorflow/compiler/mlir/lite/transforms/reduce_type_precision.cc +++ b/tensorflow/compiler/mlir/lite/transforms/reduce_type_precision.cc @@ -80,7 +80,7 @@ class CheckRangeAndConvertI8ToI4 : public OpRewritePattern { mlir::RankedTensorType::get(const_type.getShape(), builder.getI4Type()); auto newAttr = DenseElementsAttr::getFromRawBuffer( shaped_type, mlir::cast(op.getValue()).getRawData()); - rewriter.replaceOpWithNewOp(op, newAttr); + rewriter.replaceOpWithNewOp(op, newAttr); return success(); } @@ -108,7 +108,7 @@ class SanitizeGatherOpOutputToI4 : public OpRewritePattern { Builder builder(op.getContext()); auto new_gather_op = TFL::GatherOp::create(rewriter, op.getLoc(), - /*result=*/ + /*resultTypes=*/ mlir::cast(op.getResult().getType()) .clone(builder.getI4Type()), /*operands=*/op.getOperands(), op->getAttrs()); diff --git a/tensorflow/core/kernels/gather_nd_op.h b/tensorflow/core/kernels/gather_nd_op.h index 352289b4c554d0..1295b5e8c5b7f0 100644 --- a/tensorflow/core/kernels/gather_nd_op.h +++ b/tensorflow/core/kernels/gather_nd_op.h @@ -156,7 +156,7 @@ absl::Status DoGatherNd( using CPUDevice = Eigen::ThreadPoolDevice; const bool check_bad_indices = - ((std::is_same::value && + ((std::is_same_v && bad_indices_policy == BadIndicesPolicy::kDefault) || bad_indices_policy == BadIndicesPolicy::kError); if (check_bad_indices && bad_i >= 0) { diff --git a/tensorflow/lite/core/api/flatbuffer_conversions.cc b/tensorflow/lite/core/api/flatbuffer_conversions.cc index d47767c79b0c3b..f0d93ff9d2bfc6 100644 --- a/tensorflow/lite/core/api/flatbuffer_conversions.cc +++ b/tensorflow/lite/core/api/flatbuffer_conversions.cc @@ -77,12 +77,12 @@ void CheckParsePointerParams(const Operator* op, ErrorReporter* error_reporter, TFLITE_DCHECK(builtin_data != nullptr); } -// Copies the contents from the flatbuffer int vector `flatbuffer` into the +// Copies the contents from the flatbuffer int vector `flat_vector` into the // int array `buffer`. `flat_vector` and `buffer` represent the same // configuration operation for a given operation. template static TfLiteStatus FlatBufferIntVectorToArray( - int max_size_of_buffer, const flatbuffers::Vector* flat_vector, + size_t max_size_of_buffer, const flatbuffers::Vector* flat_vector, DataType* buffer, ErrorReporter* error_reporter, const char* op_name) { if (!flat_vector) { TF_LITE_REPORT_ERROR(error_reporter, @@ -2248,7 +2248,7 @@ TfLiteStatus ParseStablehloScatter(const Operator* op, if (schema_params->update_window_dims()) { TF_LITE_ENSURE_STATUS(FlatBufferIntVectorToArray( - schema_params->update_window_dims()->size() * sizeof(int64_t), + sizeof(params->update_window_dims), schema_params->update_window_dims(), params->update_window_dims, error_reporter, "stablehlo_scatter")); params->num_update_window_dims = @@ -2257,7 +2257,7 @@ TfLiteStatus ParseStablehloScatter(const Operator* op, if (schema_params->inserted_window_dims()) { TF_LITE_ENSURE_STATUS(FlatBufferIntVectorToArray( - schema_params->inserted_window_dims()->size() * sizeof(int64_t), + sizeof(params->inserted_window_dims), schema_params->inserted_window_dims(), params->inserted_window_dims, error_reporter, "stablehlo_scatter")); params->num_inserted_window_dims = @@ -2266,8 +2266,7 @@ TfLiteStatus ParseStablehloScatter(const Operator* op, if (schema_params->scatter_dims_to_operand_dims()) { TF_LITE_ENSURE_STATUS(FlatBufferIntVectorToArray( - schema_params->scatter_dims_to_operand_dims()->size() * - sizeof(int64_t), + sizeof(params->scatter_dims_to_operand_dims), schema_params->scatter_dims_to_operand_dims(), params->scatter_dims_to_operand_dims, error_reporter, "stablehlo_scatter")); @@ -2332,8 +2331,7 @@ TfLiteStatus ParseStablehloGather(const Operator* op, if (schema_params != nullptr) { if (schema_params->offset_dims()) { TF_LITE_ENSURE_STATUS(FlatBufferIntVectorToArray( - /*max_size_of_buffer=*/schema_params->offset_dims()->size() * - sizeof(int64_t), + /*max_size_of_buffer=*/sizeof(params->offset_dims), /*flat_vector=*/schema_params->offset_dims(), /*buffer=*/params->offset_dims, /*error_reporter=*/error_reporter, /*op_name=*/"stablehlo_gather")); @@ -2342,7 +2340,7 @@ TfLiteStatus ParseStablehloGather(const Operator* op, if (schema_params->collapsed_slice_dims()) { TF_LITE_ENSURE_STATUS(FlatBufferIntVectorToArray( - schema_params->collapsed_slice_dims()->size() * sizeof(int64_t), + sizeof(params->collapsed_slice_dims), schema_params->collapsed_slice_dims(), params->collapsed_slice_dims, error_reporter, "stablehlo_gather")); params->num_collapsed_slice_dims = @@ -2351,9 +2349,8 @@ TfLiteStatus ParseStablehloGather(const Operator* op, if (schema_params->start_index_map()) { TF_LITE_ENSURE_STATUS(FlatBufferIntVectorToArray( - schema_params->start_index_map()->size() * sizeof(int64_t), - schema_params->start_index_map(), params->start_index_map, - error_reporter, "stablehlo_gather")); + sizeof(params->start_index_map), schema_params->start_index_map(), + params->start_index_map, error_reporter, "stablehlo_gather")); params->num_start_index_map = schema_params->start_index_map()->size(); } @@ -2361,9 +2358,8 @@ TfLiteStatus ParseStablehloGather(const Operator* op, if (schema_params->slice_sizes()) { TF_LITE_ENSURE_STATUS(FlatBufferIntVectorToArray( - schema_params->slice_sizes()->size() * sizeof(int64_t), - schema_params->slice_sizes(), params->slice_sizes, error_reporter, - "stablehlo_gather")); + sizeof(params->slice_sizes), schema_params->slice_sizes(), + params->slice_sizes, error_reporter, "stablehlo_gather")); params->num_slice_sizes = schema_params->slice_sizes()->size(); } diff --git a/tensorflow/lite/core/api/flatbuffer_conversions_test.cc b/tensorflow/lite/core/api/flatbuffer_conversions_test.cc index 0903795850a417..093022deeb3798 100644 --- a/tensorflow/lite/core/api/flatbuffer_conversions_test.cc +++ b/tensorflow/lite/core/api/flatbuffer_conversions_test.cc @@ -943,4 +943,97 @@ TEST_F(StablehloPadFlatbufferConversionsTest, DeathTests) { ""); } +class StablehloScatterFlatbufferConversionsTest + : public FlatbufferConversionsTest { + protected: + static constexpr int kMaxDims = + TFLITE_STABLEHLO_SCATTER_PARAMS_MAX_DIMENSION_COUNT; + + const Operator* BuildScatterOperator(int update_window_dims_count, + int inserted_window_dims_count, + int scatter_dims_count) { + const auto update_window_dims = builder_.CreateVector( + std::vector(update_window_dims_count, 1)); + const auto inserted_window_dims = builder_.CreateVector( + std::vector(inserted_window_dims_count, 1)); + const auto scatter_dims = + builder_.CreateVector(std::vector(scatter_dims_count, 1)); + return BuildTestOperator( + BuiltinOptions2_StablehloScatterOptions, + CreateStablehloScatterOptions( + builder_, /*indices_are_sorted=*/true, update_window_dims, + inserted_window_dims, scatter_dims, + /*index_vector_dim=*/0, /*unique_indices=*/false, + /*update_computation_subgraph_index=*/0) + .Union()); + } + + void ExpectTooManyDimensions(int update_window_dims_count, + int inserted_window_dims_count, + int scatter_dims_count) { + TfLiteStablehloScatterParams* output_data = nullptr; + EXPECT_EQ(ParseOpData(BuildScatterOperator(update_window_dims_count, + inserted_window_dims_count, + scatter_dims_count), + BuiltinOperator_STABLEHLO_SCATTER, &mock_reporter_, + &mock_allocator_, (void**)&output_data), + kTfLiteError); + EXPECT_EQ(output_data, nullptr); + EXPECT_THAT(mock_reporter_.GetString(), + HasSubstr("Found too many dimensions in the input array of " + "operation 'stablehlo_scatter'.")); + } +}; + +TEST_F(StablehloScatterFlatbufferConversionsTest, AcceptsMaximumDimensions) { + TfLiteStablehloScatterParams* output_data = nullptr; + EXPECT_EQ(ParseOpData(BuildScatterOperator(kMaxDims, kMaxDims, kMaxDims), + BuiltinOperator_STABLEHLO_SCATTER, &mock_reporter_, + &mock_allocator_, (void**)&output_data), + kTfLiteOk); + ASSERT_NE(output_data, nullptr); + EXPECT_EQ(output_data->num_update_window_dims, kMaxDims); + EXPECT_EQ(output_data->num_inserted_window_dims, kMaxDims); + EXPECT_EQ(output_data->num_scatter_dims_to_operand_dims, kMaxDims); +} + +TEST_F(StablehloScatterFlatbufferConversionsTest, + RejectsTooManyUpdateWindowDimensions) { + ExpectTooManyDimensions(kMaxDims + 1, 1, 1); +} + +TEST_F(StablehloScatterFlatbufferConversionsTest, + RejectsTooManyInsertedWindowDimensions) { + ExpectTooManyDimensions(1, kMaxDims + 1, 1); +} + +TEST_F(StablehloScatterFlatbufferConversionsTest, + RejectsTooManyScatterDimensions) { + ExpectTooManyDimensions(1, 1, kMaxDims + 1); +} + +TEST_F(FlatbufferConversionsTest, + ParseStablehloGatherRejectsTooManyDimensions) { + std::vector too_many_dims( + TFLITE_STABLEHLO_GATHER_PARAMS_MAX_DIMENSION_COUNT + 1, 1); + const Operator* op = BuildTestOperator( + BuiltinOptions2_StablehloGatherOptions, + CreateStablehloGatherOptions( + builder_, /*offset_dims=*/builder_.CreateVector(too_many_dims), + /*collapsed_slice_dims=*/builder_.CreateVector({1}), + /*start_index_map=*/builder_.CreateVector({1}), + /*index_vector_dim=*/0, + /*slice_sizes=*/builder_.CreateVector({1}), + /*indices_are_sorted=*/true) + .Union()); + void* output_data = nullptr; + EXPECT_EQ(ParseOpData(op, BuiltinOperator_STABLEHLO_GATHER, &mock_reporter_, + &mock_allocator_, &output_data), + kTfLiteError); + EXPECT_EQ(output_data, nullptr); + EXPECT_THAT(mock_reporter_.GetString(), + HasSubstr("Found too many dimensions in the input array of " + "operation 'stablehlo_gather'.")); +} + } // namespace tflite diff --git a/tensorflow/python/ops/numpy_ops/BUILD b/tensorflow/python/ops/numpy_ops/BUILD index 01fe3fe3a74d04..a9f146815d06bc 100644 --- a/tensorflow/python/ops/numpy_ops/BUILD +++ b/tensorflow/python/ops/numpy_ops/BUILD @@ -266,6 +266,18 @@ cuda_py_strict_test( ], ) +cuda_py_strict_test( + name = "np_array_ops_no_numpy_methods_test", + srcs = ["np_array_ops_no_numpy_methods_test.py"], + deps = [ + ":np_array_ops", + "//tensorflow/python/framework:constant_op", + "//tensorflow/python/framework:ops", + "//tensorflow/python/platform:client_testlib", + "//third_party/py/numpy", + ], +) + cuda_py_strict_test( name = "np_math_ops_no_numpy_methods_test", srcs = ["np_math_ops_no_numpy_methods_test.py"], diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py index 42caa5e7d514fd..85d12f133b612e 100644 --- a/tensorflow/python/ops/numpy_ops/np_array_ops.py +++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py @@ -853,12 +853,15 @@ def around(a, decimals=0): # pylint: disable=missing-docstring # Use float as the working dtype when a.dtype is exact (e.g. integer), # because `decimals` can be negative. float_dtype = np_utils.result_type(float) - a = a.astype(float_dtype) + # `Tensor.astype` only exists once `enable_numpy_methods_on_tensor()` + # has been called, and is an alias of `math_ops.cast`; calling `cast` + # directly keeps `around` working without that opt-in. + a = math_ops.cast(a, float_dtype) factor = math_ops.cast(factor, float_dtype) a = math_ops.multiply(a, factor) a = math_ops.round(a) a = math_ops.divide(a, factor) - return a.astype(dtype) + return math_ops.cast(a, dtype) setattr(np_arrays.ndarray, '__round__', around) @@ -2177,7 +2180,9 @@ def _with_index_update_helper(update_method, a, slice_spec, updates): a_dtype = a.dtype a, updates = _promote_dtype_binary(a, updates) result_t = _slice_helper(a, slice_spec, update_method, updates) - return result_t.astype(a_dtype) + # See the note in `around`: `astype` depends on an opt-in that this + # module-level helper does not require of its callers. + return math_ops.cast(result_t, a_dtype) setattr(np_arrays.ndarray, '_numpy_style_getitem', _getitem) diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_no_numpy_methods_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_no_numpy_methods_test.py new file mode 100644 index 00000000000000..803699f0503dbb --- /dev/null +++ b/tensorflow/python/ops/numpy_ops/np_array_ops_no_numpy_methods_test.py @@ -0,0 +1,114 @@ +# Copyright 2026 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for tf numpy array methods that must not depend on Tensor methods. + +`np_math_ops.enable_numpy_methods_on_tensor()` adds numpy methods such as +`astype` to `Tensor`. The other numpy_ops test modules call it from their +`__main__`, so a code path that only works once those methods are installed +still passes there. This module deliberately does not call it, which is the +configuration a user gets from a plain `import tensorflow`. +""" + +import numpy as np + +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import ops +from tensorflow.python.ops.numpy_ops import np_array_ops +from tensorflow.python.platform import test + + +class ArrayWithoutNumpyMethodsOnTensorTest(test.TestCase): + + def testAroundAcceptsFloatInputs(self): + # `around` casts back to the argument's dtype on every path, so it is + # broken for float arguments too, not only the promoted integer ones. + arg = np.array([1.234, 5.678, -2.345], dtype=np.float32) + for decimals in (0, 1, 2): + actual = np_array_ops.around(arg, decimals) + np.testing.assert_allclose( + np.asarray(actual), + np.around(arg, decimals), + rtol=1e-6, + atol=1e-6, + err_msg='around({}, {})'.format(arg, decimals), + ) + + def testAroundAcceptsIntegerInputs(self): + # An integer argument takes the promoting branch, which computes in a + # float dtype and casts back to the integer dtype at the end. + for arg in ( + np.array([11, 25, 37], dtype=np.int32), + np.array([11, 25, 37], dtype=np.int64), + ): + for decimals in (0, 1): + actual = np_array_ops.around(arg, decimals) + np.testing.assert_array_equal( + np.asarray(actual), + np.around(arg, decimals), + err_msg='around({}, {})'.format(arg, decimals), + ) + + def testAroundPreservesDtype(self): + for dtype in (np.int32, np.int64, np.float32, np.float64): + arg = np.array([1, 2, 3], dtype=dtype) + self.assertEqual(np_array_ops.around(arg).dtype, dtype) + + def testRoundAcceptsFloatInputs(self): + arg = np.array([1.234, 5.678], dtype=np.float32) + np.testing.assert_allclose( + np.asarray(np_array_ops.round(arg, 1)), + np.round(arg, 1), + rtol=1e-6, + atol=1e-6, + ) + + def testBuiltinRoundOnTensor(self): + # `around` is installed as `Tensor.__round__` at import time, so the + # builtin `round()` reaches the same code path. + tensor = constant_op.constant([1.234, 5.678], dtype='float32') + np.testing.assert_allclose( + np.asarray(round(tensor, 1)), + np.around(np.array([1.234, 5.678], dtype=np.float32), 1), + rtol=1e-6, + atol=1e-6, + ) + + def testIndexUpdateHelperPreservesDtype(self): + # `_with_index_update_helper` backs the `_with_index_*` methods, which are + # attached to `Tensor` at import time rather than by + # `enable_numpy_methods_on_tensor()`, so it must not require the opt-in + # either. It is called directly here because the methods are attached as + # `functools.partial` objects, and whether those bind the tensor as their + # first argument differs between Python versions. + tensor = constant_op.constant([1, 2, 3, 4], dtype='int32') + updates = constant_op.constant([9, 9], dtype='int32') + cases = [ + (np_array_ops._UpdateMethod.UPDATE, np.array([9, 9, 3, 4])), + (np_array_ops._UpdateMethod.ADD, np.array([10, 11, 3, 4])), + ] + for update_method, expected in cases: + actual = np_array_ops._with_index_update_helper( + update_method, tensor, slice(0, 2), updates + ) + self.assertEqual(actual.dtype, tensor.dtype, msg=str(update_method)) + np.testing.assert_array_equal( + np.asarray(actual), expected, err_msg=str(update_method) + ) + + +if __name__ == '__main__': + ops.enable_eager_execution() + # Intentionally not calling `np_math_ops.enable_numpy_methods_on_tensor()`. + test.main() diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py index b31cc630874afb..18e885f5dbdc0d 100644 --- a/tensorflow/python/ops/numpy_ops/np_math_ops.py +++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py @@ -443,7 +443,9 @@ def f(x1, x2): y = _bin_op(f, x1, x2) if not np.issubdtype(y.dtype.as_numpy_dtype, np.inexact): - y = y.astype(np_utils.result_type(float)) + # See the note in `_scalar`: `astype` is unavailable without the + # `enable_numpy_methods_on_tensor()` opt-in. + y = math_ops.cast(y, np_utils.result_type(float)) return y @@ -1557,7 +1559,9 @@ def average(a, axis=None, weights=None, returned=False): # pylint: disable=miss default_float_type = np_utils.result_type(float) if weights is None: # Treat all weights as 1 if not np.issubdtype(a.dtype.as_numpy_dtype, np.inexact): - a = a.astype(np_utils.result_type(a.dtype, default_float_type)) + # See the note in `_scalar`: `astype` is unavailable without the + # `enable_numpy_methods_on_tensor()` opt-in. + a = math_ops.cast(a, np_utils.result_type(a.dtype, default_float_type)) avg = math_ops.reduce_mean(a, axis=axis) if returned: if axis is None: diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops_no_numpy_methods_test.py b/tensorflow/python/ops/numpy_ops/np_math_ops_no_numpy_methods_test.py index 73833e41036f04..18c50d31fbdbce 100644 --- a/tensorflow/python/ops/numpy_ops/np_math_ops_no_numpy_methods_test.py +++ b/tensorflow/python/ops/numpy_ops/np_math_ops_no_numpy_methods_test.py @@ -95,6 +95,37 @@ def testUnaryOpsAcceptIntegerScalars(self): err_msg='{}(2)'.format(name), ) + def testHeavisideAcceptsIntegerInputs(self): + # `heaviside` promotes an exact result to a float dtype, which must not + # depend on `enable_numpy_methods_on_tensor()`. + x1 = np.array([-1, 0, 1], dtype=np.int32) + for x2 in (np.int32(1), np.array([5, 5, 5], dtype=np.int32)): + actual = np_math_ops.heaviside(x1, x2) + np.testing.assert_allclose( + np.asarray(actual), + np.heaviside(x1, x2), + rtol=1e-6, + atol=1e-6, + err_msg='heaviside({}, {})'.format(x1, x2), + ) + + def testAverageAcceptsIntegerInputs(self): + # The unweighted branch of `average` promotes an integer argument to a + # float dtype before reducing. + for arg in ( + [1, 2, 3], + np.array([1, 2, 3], dtype=np.int32), + np.array([1, 2, 3], dtype=np.int64), + ): + actual = np_math_ops.average(arg) + np.testing.assert_allclose( + np.asarray(actual), + np.average(np.asarray(arg)), + rtol=1e-6, + atol=1e-6, + err_msg='average({})'.format(arg), + ) + def testFloatInputsAreUnchanged(self): arg = np.array([1.5, 2.5, 3.5], dtype=np.float32) for name, tf_fun, np_fun in _PROMOTING_UNARY_OPS: diff --git a/third_party/xla/build_tools/lint/tags.py b/third_party/xla/build_tools/lint/tags.py index ffb5c2bbd74707..a909af93153f8e 100644 --- a/third_party/xla/build_tools/lint/tags.py +++ b/third_party/xla/build_tools/lint/tags.py @@ -88,6 +88,7 @@ "requires-gpu-sm90-only": "Requires exactly sm90.", "requires-gpu-sm100-only": "Requires exactly sm100.", "requires-gpu-sm103-only": "Requires exactly sm103.", + "requires-gpu-sm107-only": "Requires exactly sm107.", "requires-gpu-sm120-only": "Requires exactly sm120.", "full": ( "Test requires a full GPU, not a partitioned one. No effect in OSS." @@ -116,6 +117,7 @@ "xla_gb200": "Runs on a gb200.", "xla_gb300": "Runs on a gb300.", "xla_rtx6000pro": "Runs on an rtx6000pro.", + "xla_vr200": "Runs on a vr200.", "xla_device_p100": "Runs on a p100.", "xla_device_v100": "Runs on a v100.", "xla_device_a100": "Runs on an a100.", @@ -124,6 +126,7 @@ "xla_device_gb200": "Runs on a gb200.", "xla_device_gb300": "Runs on a gb300.", "xla_device_rtx6000pro": "Runs on an rtx6000pro.", + "xla_device_vr200": "Runs on a vr200.", # Below tags are consumed by `xla_test`. "test_migrated_to_hlo_runner_pjrt": ( "Adds the appropriate `xla/tests:pjrt_$BACKEND_client_registry` to the" diff --git a/third_party/xla/xla/BUILD b/third_party/xla/xla/BUILD index 0325cc3a0d4984..fc24f6ecb62a80 100644 --- a/third_party/xla/xla/BUILD +++ b/third_party/xla/xla/BUILD @@ -1311,6 +1311,7 @@ cc_library( "//xla/tsl/util:command_line_flags", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base", + "@com_google_absl//absl/base:config", "@com_google_absl//absl/base:no_destructor", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", diff --git a/third_party/xla/xla/backends/gpu/tests/dynamic_slice_fusion_v2_test.cc b/third_party/xla/xla/backends/gpu/tests/dynamic_slice_fusion_v2_test.cc index 7ebcc63fd453b5..27bfad7ff755d7 100644 --- a/third_party/xla/xla/backends/gpu/tests/dynamic_slice_fusion_v2_test.cc +++ b/third_party/xla/xla/backends/gpu/tests/dynamic_slice_fusion_v2_test.cc @@ -15,6 +15,7 @@ limitations under the License. #include #include +#include #include #include diff --git a/third_party/xla/xla/backends/gpu/tests/gpu_unrolling_test.cc b/third_party/xla/xla/backends/gpu/tests/gpu_unrolling_test.cc index 23f399566ea852..0ad31a3f932563 100644 --- a/third_party/xla/xla/backends/gpu/tests/gpu_unrolling_test.cc +++ b/third_party/xla/xla/backends/gpu/tests/gpu_unrolling_test.cc @@ -13,6 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include #include #include diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc index 66c478dd3879da..73534629847b27 100644 --- a/third_party/xla/xla/debug_options_flags.cc +++ b/third_party/xla/xla/debug_options_flags.cc @@ -30,6 +30,7 @@ limitations under the License. #include "absl/algorithm/container.h" #include "absl/base/call_once.h" +#include "absl/base/config.h" // IWYU pragma: keep #include "absl/base/no_destructor.h" #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" @@ -564,9 +565,19 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() { opts.set_xla_gpu_native_emitter_tune_unroll_factor_for_loops(false); opts.set_xla_gpu_experimental_use_ragged_dot_fusion(false); - opts.set_xla_cpu_collective_call_warn_stuck_seconds(20); - opts.set_xla_cpu_collective_call_terminate_timeout_seconds(40); - opts.set_xla_cpu_collective_timeout_seconds(30 * 60); +#if defined(ABSL_HAVE_ADDRESS_SANITIZER) || \ + defined(ABSL_HAVE_HWADDRESS_SANITIZER) || \ + defined(ABSL_HAVE_MEMORY_SANITIZER) || defined(ABSL_HAVE_THREAD_SANITIZER) + constexpr int kSanitizerMultiplier = 10; +#else + constexpr int kSanitizerMultiplier = 1; +#endif + + opts.set_xla_cpu_collective_call_warn_stuck_seconds(20 * + kSanitizerMultiplier); + opts.set_xla_cpu_collective_call_terminate_timeout_seconds( + 40 * kSanitizerMultiplier); + opts.set_xla_cpu_collective_timeout_seconds(30 * 60 * kSanitizerMultiplier); opts.set_xla_keep_shardings_after_spmd(false); opts.set_xla_enable_hlo_sharding_v3(false); diff --git a/third_party/xla/xla/tests/backend_defs.bzl b/third_party/xla/xla/tests/backend_defs.bzl index 8b2eb3df05510b..e6ee22a9bf6226 100644 --- a/third_party/xla/xla/tests/backend_defs.bzl +++ b/third_party/xla/xla/tests/backend_defs.bzl @@ -34,6 +34,7 @@ NVIDIA_GPU_BACKENDS = [ "b200", "gb200", "gb300", + "vr200", ] + if_google([], ["rtx6000pro"]) # The generic "gpu" backend includes the actual backends in this list. @@ -44,6 +45,7 @@ NVIDIA_GPU_DEFAULT_BACKENDS = [ "b200", "gb200", "gb300", + "vr200", ] + if_google([], ["rtx6000pro"]) AMD_GPU_DEFAULT_BACKENDS = ["amdgpu_any"] @@ -93,6 +95,7 @@ def prepare_nvidia_gpu_backend_data(backends, disabled_backends, backend_tags, b "b200": (10, 0), "gb200": (10, 0), "gb300": (10, 3), + "vr200": (10, 7), "rtx6000pro": (12, 0), } for gpu_backend in NVIDIA_GPU_BACKENDS: diff --git a/third_party/xla/xla/tests/constraint_propagator.cc b/third_party/xla/xla/tests/constraint_propagator.cc index 8ce31e5b30cff4..15399c6d887e8f 100644 --- a/third_party/xla/xla/tests/constraint_propagator.cc +++ b/third_party/xla/xla/tests/constraint_propagator.cc @@ -1001,11 +1001,11 @@ absl::Status ConstraintPropagator::PropagateFusionBoundary( const HloInstruction* fused_root = fusion_instruction->fused_expression_root(); if (fused_root != nullptr) { - // Backward: outer constraint on fusion result flows into inner root. ConstraintState fusion_state = states_[fusion_instruction]; + ConstraintState root_state = states_[fused_root]; + // Backward: outer constraint on fusion result flows into inner root. states_[fused_root].Merge(fusion_state); // Forward: internal constraint computed on root flows out to fusion result. - ConstraintState root_state = states_[fused_root]; states_[fusion_instruction].Merge(root_state); } @@ -1016,13 +1016,13 @@ absl::Status ConstraintPropagator::PropagateFusionBoundary( if (fused_param == nullptr) { continue; } + ConstraintState operand_state = states_[operand]; + ConstraintState param_state = states_[fused_param]; // Backward: constraints accumulated on the internal parameter flow out // to the caller operand. - ConstraintState param_state = states_[fused_param]; states_[operand].Merge(param_state); // Forward: constraints established on the caller operand flow into the // internal parameter. - ConstraintState operand_state = states_[operand]; states_[fused_param].Merge(operand_state); }