[pull] master from tensorflow:master

tensorflow/compiler/mlir/lite/BUILD

@@ -1511,7 +1511,7 @@ cc_library(
         "transforms/prepare_quantize_helper.cc",
         "transforms/quantization/bias_quantizer_pass.cc",
         "transforms/quantization/fuse_qdq_pass.cc",
-        "transforms/quantization/propagate_qsv_pass.cc",
+        "transforms/quantization/propagate_qparams_pass.cc",
         "transforms/quantization/quant_utils.cc",
         "transforms/quantize.cc",
         "transforms/quantize_variables.cc",

tensorflow/compiler/mlir/lite/stablehlo/BUILD

@@ -76,6 +76,35 @@ cc_library(
     alwayslink = 1,
 )
 
+cc_library(
+    name = "legalize_vhlo_quant_custom_calls",
+    srcs = [
+        "transforms/legalize_vhlo_quant_custom_calls.cc",
+    ],
+    hdrs = [
+        "transforms/stablehlo_passes.h",
+        "transforms/stablehlo_passes.h.inc",
+    ],
+    copts = [
+        "-Ithird_party",
+    ],
+    deps = [
+        ":passes_inc_gen",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+        "@llvm-project//mlir:Transforms",
+        "@stablehlo//:stablehlo_ops",
+        "@stablehlo//:vhlo_ops",
+        "@stablehlo//:vhlo_types",
+    ],
+    alwayslink = 1,
+)
+
 cc_library(
     name = "stablehlo_util",
     srcs = [
@@ -355,7 +384,6 @@ cc_library(
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:QuantOps",
         "@llvm-project//mlir:Support",
-        "@llvm-project//mlir:Transforms",
         "@xla//xla/mlir_hlo",
     ],
     alwayslink = 1,
@@ -405,6 +433,7 @@ cc_library(
         "-Ithird_party",
     ],
     deps = [
+        ":legalize_stablehlo_custom_call_to_composite",
         "//tensorflow/compiler/mlir/lite:tensorflow_lite",
         "@flatbuffers",
         "@llvm-project//llvm:Support",
@@ -937,6 +966,7 @@ tf_cc_binary(
         ":legalize_stablehlo_custom_call_to_composite",
         ":legalize_stablehlo_to_vhlo_pass",
         ":legalize_tf_xla_call_module_to_stablehlo_pass",
+        ":legalize_vhlo_quant_custom_calls",
         ":optimize",
         ":passes_inc_gen",
         ":prepare_hlo",

tensorflow/compiler/mlir/lite/stablehlo/tests/legalize-vhlo-quant-custom-calls.mlir

@@ -0,0 +1,92 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: odml-to-stablehlo-opt %s -legalize-vhlo-quant-custom-calls -split-input-file | FileCheck %s
+
+// CHECK-LABEL: func.func @legalize_vhlo_quant_dequantize
+func.func @legalize_vhlo_quant_dequantize(%arg0: tensor<4xi8>, %arg1: tensor<f32>, %arg2: tensor<i32>) -> tensor<4xf32> {
+  // CHECK-NOT: vhlo.custom_call_v1
+  // CHECK: %[[RES:.*]] = stablehlo.custom_call @quant.dequantize(%arg0, %arg1, %arg2) {axis = 0 : i32} : (tensor<4xi8>, tensor<f32>, tensor<i32>) -> tensor<4xf32>
+  // CHECK: return %[[RES]] : tensor<4xf32>
+  %0 = "vhlo.custom_call_v1"(%arg0, %arg1, %arg2) <{
+    api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
+    backend_config = #vhlo.string_v1<"">,
+    call_target_name = #vhlo.string_v1<"quant.dequantize">,
+    called_computations = #vhlo.array_v1<[]>,
+    has_side_effect = #vhlo.bool_v1<false>,
+    operand_layouts = #vhlo.array_v1<[]>,
+    output_operand_aliases = #vhlo.array_v1<[]>,
+    result_layouts = #vhlo.array_v1<[]>
+  }> {axis = #vhlo.integer_v1<0 : i32>} : (tensor<4xi8>, tensor<f32>, tensor<i32>) -> tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @legalize_vhlo_quant_quantize
+func.func @legalize_vhlo_quant_quantize(%arg0: tensor<4xf32>, %arg1: tensor<f32>, %arg2: tensor<i32>) -> tensor<4xi8> {
+  // CHECK-NOT: vhlo.custom_call_v1
+  // CHECK: %[[RES:.*]] = stablehlo.custom_call @quant.quantize(%arg0, %arg1, %arg2) : (tensor<4xf32>, tensor<f32>, tensor<i32>) -> tensor<4xi8>
+  // CHECK: return %[[RES]] : tensor<4xi8>
+  %0 = "vhlo.custom_call_v1"(%arg0, %arg1, %arg2) <{
+    api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
+    backend_config = #vhlo.string_v1<"">,
+    call_target_name = #vhlo.string_v1<"quant.quantize">,
+    called_computations = #vhlo.array_v1<[]>,
+    has_side_effect = #vhlo.bool_v1<false>,
+    operand_layouts = #vhlo.array_v1<[]>,
+    output_operand_aliases = #vhlo.array_v1<[]>,
+    result_layouts = #vhlo.array_v1<[]>
+  }> : (tensor<4xf32>, tensor<f32>, tensor<i32>) -> tensor<4xi8>
+  return %0 : tensor<4xi8>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @legalize_vhlo_quant_fake_quant
+func.func @legalize_vhlo_quant_fake_quant(%arg0: tensor<4xf32>, %arg1: tensor<f32>, %arg2: tensor<i32>) -> tensor<4xf32> {
+  // CHECK-NOT: vhlo.custom_call_v1
+  // CHECK: %[[RES:.*]] = stablehlo.custom_call @quant.fake_quant(%arg0, %arg1, %arg2) {narrow_range = false} : (tensor<4xf32>, tensor<f32>, tensor<i32>) -> tensor<4xf32>
+  // CHECK: return %[[RES]] : tensor<4xf32>
+  %0 = "vhlo.custom_call_v1"(%arg0, %arg1, %arg2) <{
+    api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
+    backend_config = #vhlo.string_v1<"">,
+    call_target_name = #vhlo.string_v1<"quant.fake_quant">,
+    called_computations = #vhlo.array_v1<[]>,
+    has_side_effect = #vhlo.bool_v1<false>,
+    operand_layouts = #vhlo.array_v1<[]>,
+    output_operand_aliases = #vhlo.array_v1<[]>,
+    result_layouts = #vhlo.array_v1<[]>
+  }> {narrow_range = #vhlo.bool_v1<false>} : (tensor<4xf32>, tensor<f32>, tensor<i32>) -> tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @keep_other_vhlo_custom_call
+func.func @keep_other_vhlo_custom_call(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+  // CHECK: "vhlo.custom_call_v1"(%arg0)
+  // CHECK-SAME: call_target_name = #vhlo.string_v1<"other.custom_call">
+  %0 = "vhlo.custom_call_v1"(%arg0) <{
+    api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
+    backend_config = #vhlo.string_v1<"">,
+    call_target_name = #vhlo.string_v1<"other.custom_call">,
+    called_computations = #vhlo.array_v1<[]>,
+    has_side_effect = #vhlo.bool_v1<false>,
+    operand_layouts = #vhlo.array_v1<[]>,
+    output_operand_aliases = #vhlo.array_v1<[]>,
+    result_layouts = #vhlo.array_v1<[]>
+  }> : (tensor<4xf32>) -> tensor<4xf32>
+  return %0 : tensor<4xf32>
+}

tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_vhlo_quant_custom_calls.cc

@@ -0,0 +1,218 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <memory>
+#include <utility>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo  // build_cleaner: keep
+#include "stablehlo/dialect/VhloOps.h"  // from @stablehlo  // build_cleaner: keep
+#include "stablehlo/dialect/VhloTypes.h"  // from @stablehlo
+#include "tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.h"
+
+namespace mlir::odml {
+
+#define GEN_PASS_DEF_LEGALIZEVHLOQUANTCUSTOMCALLSPASS
+#include "tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.h.inc"
+
+namespace {
+
+class VhloToStablehloTypeConverter : public vhlo::VhloTypeConverter {
+ public:
+  VhloToStablehloTypeConverter()
+      : vhlo::VhloTypeConverter(/*allowOtherDialects=*/true) {
+    addConversion([](Type type) -> Type { return type; });
+    addConversion([](vhlo::TokenV1Type token) -> Type {
+      return stablehlo::TokenType::get(token.getContext());
+    });
+    addVhloToBuiltinConversions();
+  }
+
+  Attribute convertEncoding(Attribute attr) const final {
+    if (auto vhloAttr =
+            mlir::dyn_cast_or_null<vhlo::TypeExtensionsV1Attr>(attr)) {
+      return stablehlo::TypeExtensionsAttr::get(vhloAttr.getContext(),
+                                                vhloAttr.getBounds());
+    }
+    return attr;
+  }
+};
+
+Attribute ConvertVhloAttrToBuiltin(Attribute attr) {
+  if (!attr) return {};
+  if (auto vhlo_str = mlir::dyn_cast<vhlo::StringV1Attr>(attr)) {
+    return StringAttr::get(attr.getContext(), vhlo_str.getValue());
+  }
+  if (auto vhlo_bool = mlir::dyn_cast<vhlo::BooleanV1Attr>(attr)) {
+    return BoolAttr::get(attr.getContext(), vhlo_bool.getValue());
+  }
+  if (auto vhlo_int = mlir::dyn_cast<vhlo::IntegerV1Attr>(attr)) {
+    VhloToStablehloTypeConverter type_converter;
+    Type type = type_converter.convertType(vhlo_int.getType());
+    if (!type) type = vhlo_int.getType();
+    return IntegerAttr::get(type, vhlo_int.getValue());
+  }
+  if (auto vhlo_float = mlir::dyn_cast<vhlo::FloatV1Attr>(attr)) {
+    VhloToStablehloTypeConverter type_converter;
+    Type type = type_converter.convertType(vhlo_float.getType());
+    if (!type) type = vhlo_float.getType();
+    return FloatAttr::get(type, vhlo_float.getValue());
+  }
+  if (auto vhlo_tensor = mlir::dyn_cast<vhlo::TensorV1Attr>(attr)) {
+    return vhlo_tensor.getData();
+  }
+  return attr;
+}
+
+struct LegalizeVhloQuantCustomCallPattern : public RewritePattern {
+  explicit LegalizeVhloQuantCustomCallPattern(MLIRContext* context)
+      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}
+
+  LogicalResult matchAndRewrite(Operation* op,
+                                PatternRewriter& rewriter) const override {
+    StringRef op_name = op->getName().getStringRef();
+    if (op_name != "vhlo.custom_call_v1" && op_name != "vhlo.custom_call") {
+      return failure();
+    }
+
+    Attribute raw_target = op->getAttr("call_target_name");
+    if (!raw_target) return failure();
+
+    StringRef target_name;
+    if (auto attr = mlir::dyn_cast<vhlo::StringV1Attr>(raw_target)) {
+      target_name = attr.getValue();
+    } else if (auto attr = mlir::dyn_cast<StringAttr>(raw_target)) {
+      target_name = attr.getValue();
+    }
+
+    if (target_name != "quant.dequantize" && target_name != "quant.quantize" &&
+        target_name != "quant.fake_quant") {
+      return failure();
+    }
+
+    VhloToStablehloTypeConverter type_converter;
+
+    SmallVector<Type> result_types;
+    for (Type t : op->getResultTypes()) {
+      Type conv = type_converter.convertType(t);
+      result_types.push_back(conv ? conv : t);
+    }
+
+    SmallVector<Value> operands;
+    for (Value val : op->getOperands()) {
+      Type conv = type_converter.convertType(val.getType());
+      if (conv && conv != val.getType()) {
+        val =
+            rewriter.create<UnrealizedConversionCastOp>(op->getLoc(), conv, val)
+                .getResult(0);
+      }
+      operands.push_back(val);
+    }
+
+    SmallVector<NamedAttribute> new_attrs;
+    new_attrs.push_back(rewriter.getNamedAttr(
+        "call_target_name", rewriter.getStringAttr(target_name)));
+
+    static const char* const kIntrinsicAttrs[] = {"api_version",
+                                                  "backend_config",
+                                                  "call_target_name",
+                                                  "called_computations",
+                                                  "has_side_effect",
+                                                  "operand_layouts",
+                                                  "output_operand_aliases",
+                                                  "result_layouts",
+                                                  "result_tilings"};
+
+    for (NamedAttribute attr : op->getAttrs()) {
+      StringRef name = attr.getName().strref();
+      bool is_intrinsic = false;
+      for (const char* kAttr : kIntrinsicAttrs) {
+        if (name == kAttr) {
+          is_intrinsic = true;
+          break;
+        }
+      }
+      if (is_intrinsic) continue;
+
+      Attribute builtin_val = ConvertVhloAttrToBuiltin(attr.getValue());
+      if (builtin_val) {
+        new_attrs.push_back(rewriter.getNamedAttr(name, builtin_val));
+      }
+    }
+
+    auto new_op = rewriter.create<stablehlo::CustomCallOp>(
+        op->getLoc(), result_types, operands, new_attrs);
+
+    if (new_op->getNumResults() != op->getNumResults()) {
+      return failure();
+    }
+
+    SmallVector<Value> replacement_vals;
+    for (unsigned i = 0; i < op->getNumResults(); ++i) {
+      Value new_res = new_op.getResult(i);
+      Type orig_type = op->getResult(i).getType();
+      if (new_res.getType() != orig_type) {
+        new_res = rewriter
+                      .create<UnrealizedConversionCastOp>(op->getLoc(),
+                                                          orig_type, new_res)
+                      .getResult(0);
+      }
+      replacement_vals.push_back(new_res);
+    }
+
+    rewriter.replaceOp(op, replacement_vals);
+    return success();
+  }
+};
+
+class LegalizeVhloQuantCustomCallsPass
+    : public impl::LegalizeVhloQuantCustomCallsPassBase<
+          LegalizeVhloQuantCustomCallsPass> {
+ public:
+  void runOnOperation() override {
+    ModuleOp module = getOperation();
+    MLIRContext* context = &getContext();
+    RewritePatternSet patterns(context);
+    patterns.add<LegalizeVhloQuantCustomCallPattern>(context);
+    if (failed(applyPatternsGreedily(module, std::move(patterns)))) {
+      signalPassFailure();
+    }
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<OperationPass<ModuleOp>>
+CreateLegalizeVhloQuantCustomCallsPass() {
+  return std::make_unique<LegalizeVhloQuantCustomCallsPass>();
+}
+
+static PassRegistration<LegalizeVhloQuantCustomCallsPass> pass;
+
+}  // namespace mlir::odml

tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.h

@@ -67,6 +67,10 @@ void PopulateLegalizeHloToTfPatterns(RewritePatternSet* patterns,
 // Drops vhlo/stablehlo custom calls targeting 'shape_assertion'.
 std::unique_ptr<OperationPass<ModuleOp>> CreateDropShapeAssertionsPass();
 
+// Legalizes vhlo custom calls for quantization ops to stablehlo.custom_call.
+std::unique_ptr<OperationPass<ModuleOp>>
+CreateLegalizeVhloQuantCustomCallsPass();
+
 #define GEN_PASS_DECL
 #include "tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.h.inc"