Add the skeleton for PropagateQsvPass in the tflite converter.

majiddadashi · tensorflower-gardener · commit 609397032f82 · 2025-07-24T14:17:52.000-07:00
This pass will infer QSV from surrounding tensors and in cases modifies the graph to comply with the quantization constraints of the tflite kernels.

PiperOrigin-RevId: 786055945
diff --git a/tensorflow/compiler/mlir/lite/BUILD b/tensorflow/compiler/mlir/lite/BUILD
@@ -1354,6 +1354,7 @@ cc_library(
         "transforms/prepare_quantize.cc",
         "transforms/prepare_quantize_dynamic_range.cc",
         "transforms/prepare_quantize_helper.cc",
+        "transforms/quantization/propagate_qsv_pass.cc",
         "transforms/quantize.cc",
         "transforms/quantize_variables.cc",
         "utils/generated_op_quant_spec_getters.inc",
diff --git a/tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization.td b/tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization.td
@@ -137,7 +137,7 @@ def SameOperandsAndResultsScale : OpInterface<"SameScalesOpInterface"> {
   ];
 
   let verify = [{
-    return TFL::VerifySameScales($_op);
+    return mlir::TFL::VerifySameScales($_op);
   }];
 }
 
diff --git a/tensorflow/compiler/mlir/lite/transforms/passes.h b/tensorflow/compiler/mlir/lite/transforms/passes.h
@@ -120,6 +120,10 @@ std::unique_ptr<OperationPass<func::FuncOp>> CreateDefaultQuantizePass();
 
 std::unique_ptr<OperationPass<ModuleOp>> CreateLowerQuantAnnotationsPass();
 
+// Creates an instance of the TFLite PropagateQsv pass which propagates scale
+// and zero point (QSV) information through the graph.
+std::unique_ptr<OperationPass<ModuleOp>> CreatePropagateQsvPass();
+
 // Overloading of CreateQuantizePass which takes only necessary flags to reduce
 // the binary size.
 std::unique_ptr<OperationPass<func::FuncOp>> CreateQuantizePass(
diff --git a/tensorflow/compiler/mlir/lite/transforms/passes.td b/tensorflow/compiler/mlir/lite/transforms/passes.td
@@ -354,6 +354,19 @@ def LowerQuantAnnotationsPass : Pass<"tfl-lower-quant-annotations", "mlir::Modul
   ];
 }
 
+def PropagateQsvPass : Pass<"tfl-propagate-qsv", "mlir::ModuleOp"> {
+  let summary = "Propagates Quantization Scale/Value (QSV) information through the graph.";
+  let description = [{
+      This transformation pass propagates the QSV data across operations in the
+      TensorFlow Lite dialect.
+  }];
+  let constructor = "CreatePropagateQsvPass()";
+  let dependentDialects = [
+    "TFL::TensorFlowLiteDialect",
+    "mlir::quant::QuantDialect"
+  ];
+}
+
 def QuantizeVariablesPass : Pass<"tfl-quantize-variables", "mlir::ModuleOp"> {
   let summary = "Quantize variables";
   let constructor = "CreatePrepareQuantizeVariablesPass()";
diff --git a/tensorflow/compiler/mlir/lite/transforms/quantization/propagate_qsv_pass.cc b/tensorflow/compiler/mlir/lite/transforms/quantization/propagate_qsv_pass.cc
@@ -0,0 +1,85 @@
+/* Copyright 2025 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.
+==============================================================================*/
+
+// This transformation pass propagates QSV information through the model.
+
+#include <memory>
+#include <utility>
+
+#include "mlir/Dialect/Quant/IR/Quant.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/Diagnostics.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.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/DialectConversion.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/transforms/passes.h"  // IWYU pragma: keep
+
+namespace mlir {
+namespace TFL {
+namespace {
+
+#define GEN_PASS_DEF_PROPAGATEQSVPASS
+#include "tensorflow/compiler/mlir/lite/transforms/passes.h.inc"
+
+//===----------------------------------------------------------------------===//
+// Pass Definition
+//===----------------------------------------------------------------------===//
+
+struct PropagateQsvPass : public impl::PropagateQsvPassBase<PropagateQsvPass> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PropagateQsvPass)
+
+  void runOnOperation() override;
+};
+
+//===----------------------------------------------------------------------===//
+// Pass Implementation
+//===----------------------------------------------------------------------===//
+
+void PropagateQsvPass::runOnOperation() {
+  MLIRContext* ctx = &getContext();
+  mlir::ModuleOp module = getOperation();
+
+  RewritePatternSet patterns(ctx);
+
+  // Configure the greedy pattern rewrite driver.
+  GreedyRewriteConfig greedy_config;
+  greedy_config.enableFolding(false);
+
+  // Apply the patterns.
+  if (failed(
+          applyPatternsGreedily(module, std::move(patterns), greedy_config))) {
+    signalPassFailure();
+  }
+}
+
+}  // namespace
+
+//===----------------------------------------------------------------------===//
+// Pass Creation Function
+//===----------------------------------------------------------------------===//
+
+std::unique_ptr<OperationPass<mlir::ModuleOp>> CreatePropagateQsvPass() {
+  return std::make_unique<PropagateQsvPass>();
+}
+
+}  // namespace TFL
+}  // namespace mlir

Original file line number	Diff line number	Diff line change
`@@ -137,7 +137,7 @@ def SameOperandsAndResultsScale : OpInterface<"SameScalesOpInterface"> {`
`137`	`137`	`];`
`138`	`138`
`139`	`139`	`let verify = [{`
`140`		`- return TFL::VerifySameScales($_op);`
	`140`	`+ return mlir::TFL::VerifySameScales($_op);`
`141`	`141`	`}];`
`142`	`142`	`}`
`143`	`143`