chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/lite/quantization/ir/.clang-format

@@ -0,0 +1,2 @@
+BasedOnStyle: LLVM
+AlwaysBreakTemplateDeclarations: Yes

tensorflow/compiler/mlir/lite/quantization/ir/.clang-tidy

@@ -0,0 +1,76 @@
+Checks: >
+        # LINT.IfChange(checks)
+        -*,clang-diagnostic-*,llvm-*,misc-*,
+        -misc-unused-parameters,
+        -misc-non-private-member-variables-in-classes,
+        -misc-no-recursion,
+        bugprone-argument-comment,
+        bugprone-assert-side-effect,
+        bugprone-branch-clone,
+        bugprone-copy-constructor-init,
+        bugprone-dangling-handle,
+        bugprone-dynamic-static-initializers,
+        bugprone-macro-parentheses,
+        bugprone-macro-repeated-side-effects,
+        bugprone-misplaced-widening-cast,
+        bugprone-move-forwarding-reference,
+        bugprone-multiple-statement-macro,
+        bugprone-suspicious-semicolon,
+        bugprone-swapped-arguments,
+        bugprone-terminating-continue,
+        bugprone-unused-raii,
+        bugprone-unused-return-value,
+        misc-redundant-expression,
+        misc-static-assert,
+        misc-unused-parameters,
+        misc-unused-using-decls,
+        modernize-use-bool-literals,
+        modernize-loop-convert,
+        modernize-make-unique,
+        modernize-raw-string-literal,
+        modernize-use-equals-default,
+        modernize-use-default-member-init,
+        modernize-use-emplace,
+        modernize-use-nullptr,
+        modernize-use-override,
+        modernize-use-using,
+        performance-for-range-copy,
+        performance-implicit-conversion-in-loop,
+        performance-inefficient-algorithm,
+        performance-inefficient-vector-operation,
+        performance-move-const-arg,
+        performance-no-automatic-move,
+        performance-trivially-destructible,
+        performance-unnecessary-copy-initialization,
+        performance-unnecessary-value-param,
+        readability-avoid-const-params-in-decls,
+        readability-const-return-type,
+        readability-container-size-empty,
+        readability-identifier-naming,
+        readability-inconsistent-declaration-parameter-name,
+        readability-misleading-indentation,
+        readability-redundant-control-flow,
+        readability-simplify-boolean-expr,
+        readability-simplify-subscript-expr,
+        readability-use-anyofallof
+        # LINT.ThenChange(METADATA:checks)
+
+CheckOptions:
+  # LINT.IfChange(check-options)
+  - key:             readability-identifier-naming.ClassCase
+    value:           CamelCase
+  - key:             readability-identifier-naming.EnumCase
+    value:           CamelCase
+  - key:             readability-identifier-naming.FunctionCase
+    value:           camelBack
+  - key:             readability-identifier-naming.MemberCase
+    value:           camelBack
+  - key:             readability-identifier-naming.ParameterCase
+    value:           camelBack
+  - key:             readability-identifier-naming.UnionCase
+    value:           CamelCase
+  - key:             readability-identifier-naming.VariableCase
+    value:           camelBack
+  - key:             readability-identifier-naming.IgnoreMainLikeFunctions
+    value:           1
+  # LINT.ThenChange(METADATA:check-options)

tensorflow/compiler/mlir/lite/quantization/ir/BUILD

@@ -0,0 +1,88 @@
+load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+td_library(
+    name = "QuantizationOpsTdFiles",
+    srcs = [
+        "QuantOps.td",
+        "QuantOpsBase.td",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "@llvm-project//mlir:InferTypeOpInterfaceTdFiles",
+        "@llvm-project//mlir:OpBaseTdFiles",
+        "@llvm-project//mlir:QuantizationOpsTdFiles",
+        "@llvm-project//mlir:SideEffectInterfacesTdFiles",
+    ],
+)
+
+gentbl_cc_library(
+    name = "QuantOpsIncGen",
+    compatible_with = get_compatible_with_portable(),
+    tbl_outs = {
+        "QuantOps.h.inc": ["-gen-op-decls"],
+        "QuantOps.cc.inc": ["-gen-op-defs"],
+        "QuantOpsDialect.h.inc": [
+            "-gen-dialect-decls",
+            "-dialect=quantfork",
+        ],
+        "QuantOpsDialect.cc.inc": [
+            "-gen-dialect-defs",
+            "-dialect=quantfork",
+        ],
+    },
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "QuantOps.td",
+    deps = [":QuantizationOpsTdFiles"],
+)
+
+gentbl_cc_library(
+    name = "QuantPassIncGen",
+    compatible_with = get_compatible_with_portable(),
+    tbl_outs = {"Passes.h.inc": [
+        "-gen-pass-decls",
+        "-name=quantfork",
+    ]},
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "Passes.td",
+    deps = ["@llvm-project//mlir:PassBaseTdFiles"],
+)
+
+cc_library(
+    name = "QuantOps",
+    srcs = [
+        "ConvertConst.cc",
+        "ConvertSimQuant.cc",
+        "QuantOps.cc",
+        "QuantizeUtils.cc",
+    ],
+    hdrs = [
+        "Passes.h",
+        "QuantOps.h",
+        "QuantizeUtils.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        ":QuantOpsIncGen",
+        ":QuantPassIncGen",
+        "//tensorflow/compiler/mlir/quantization/common/ir:QuantOps",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:BytecodeOpInterface",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:InferTypeOpInterface",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:QuantOps",
+        "@llvm-project//mlir:SideEffectInterfaces",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+    ],
+)

tensorflow/compiler/mlir/lite/quantization/ir/ConvertConst.cc

@@ -0,0 +1,122 @@
+/* Copyright 2022 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/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/Matchers.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/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/quantization/ir/Passes.h"
+#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantizeUtils.h"
+
+using namespace mlir;
+using namespace mlir::quantfork;
+
+using mlir::quant::QuantizedType;
+
+namespace {
+#define GEN_PASS_DEF_QUANTCONVERTCONST
+#include "tensorflow/compiler/mlir/lite/quantization/ir/Passes.h.inc"
+
+struct ConvertConstPass : public impl::QuantConvertConstBase<ConvertConstPass> {
+  void runOnOperation() override;
+};
+
+struct QuantizedConstRewrite : public OpRewritePattern<QuantizeCastOp> {
+  using OpRewritePattern<QuantizeCastOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(QuantizeCastOp qbarrier,
+                                PatternRewriter &rewriter) const override;
+};
+
+}  // namespace
+
+/// Matches a [constant] -> [qbarrier] where the qbarrier results type is
+/// quantized and the operand type is quantizable.
+
+LogicalResult QuantizedConstRewrite::matchAndRewrite(
+    QuantizeCastOp qbarrier, PatternRewriter &rewriter) const {
+  Attribute value;
+
+  // Is the operand a constant?
+  if (!matchPattern(qbarrier.getArg(), m_Constant(&value))) {
+    return failure();
+  }
+
+  // Does the qbarrier convert to a quantized type. This will not be true
+  // if a quantized type has not yet been chosen or if the cast to an equivalent
+  // storage type is not supported.
+  Type qbarrierResultType = qbarrier.getResult().getType();
+  QuantizedType quantizedElementType =
+      QuantizedType::getQuantizedElementType(qbarrierResultType);
+  if (!quantizedElementType) {
+    return failure();
+  }
+  if (!QuantizedType::castToStorageType(qbarrierResultType)) {
+    return failure();
+  }
+
+  // Is the operand type compatible with the expressed type of the quantized
+  // type? This will not be true if the qbarrier is superfluous (converts
+  // from and to a quantized type).
+  if (!quantizedElementType.isCompatibleExpressedType(
+          qbarrier.getArg().getType())) {
+    return failure();
+  }
+
+  // Is the constant value a type expressed in a way that we support?
+  if (!mlir::isa<FloatAttr, DenseElementsAttr, SparseElementsAttr>(value)) {
+    return failure();
+  }
+
+  Type newConstValueType;
+  auto newConstValue =
+      quantizeAttr(value, quantizedElementType, newConstValueType);
+  if (!newConstValue) {
+    return failure();
+  }
+
+  // When creating the new const op, use a fused location that combines the
+  // original const and the qbarrier that led to the quantization.
+  auto fusedLoc = rewriter.getFusedLoc(
+      {qbarrier.getArg().getDefiningOp()->getLoc(), qbarrier.getLoc()});
+  auto newConstOp = arith::ConstantOp::create(
+      rewriter, fusedLoc, newConstValueType, cast<TypedAttr>(newConstValue));
+  rewriter.replaceOpWithNewOp<StorageCastOp>(qbarrier, qbarrier.getType(),
+                                             newConstOp);
+  return success();
+}
+
+void ConvertConstPass::runOnOperation() {
+  RewritePatternSet patterns(&getContext());
+  auto func = getOperation();
+  auto *context = &getContext();
+  patterns.add<QuantizedConstRewrite>(context);
+  (void)applyPatternsGreedily(func, std::move(patterns));
+}
+
+std::unique_ptr<OperationPass<func::FuncOp>>
+mlir::quantfork::createConvertConstPass() {
+  return std::make_unique<ConvertConstPass>();
+}

tensorflow/compiler/mlir/lite/quantization/ir/ConvertSimQuant.cc

@@ -0,0 +1,158 @@
+/* Copyright 2022 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 <cassert>
+#include <memory>
+#include <utility>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/Types.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/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/quantization/ir/Passes.h"
+#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/quantization/common/ir/FakeQuantSupport.h"
+#include "tensorflow/compiler/mlir/quantization/common/ir/UniformSupport.h"
+
+namespace mlir::quantfork {
+
+#define GEN_PASS_DEF_QUANTCONVERTSIMULATEDQUANT
+#include "tensorflow/compiler/mlir/lite/quantization/ir/Passes.h.inc"
+
+struct ConvertSimulatedQuantPass
+    : public impl::QuantConvertSimulatedQuantBase<ConvertSimulatedQuantPass> {
+  void runOnOperation() override;
+};
+
+/// Base class rewrites ConstFakeQuant into a qbarrier/dbarrier pair.
+template <typename ConcreteRewriteClass, typename FakeQuantOp>
+class FakeQuantRewrite : public OpRewritePattern<FakeQuantOp> {
+ public:
+  using OpRewritePattern<FakeQuantOp>::OpRewritePattern;
+
+  FakeQuantRewrite(MLIRContext *ctx, bool *hadFailure)
+      : OpRewritePattern<FakeQuantOp>(ctx), hadFailure(hadFailure) {}
+
+  LogicalResult matchAndRewrite(FakeQuantOp op,
+                                PatternRewriter &rewriter) const override {
+    // TODO: If this pattern comes up more frequently, consider adding core
+    // support for failable rewrites.
+    if (failableRewrite(op, rewriter)) {
+      *hadFailure = true;
+      return failure();
+    }
+    return success();
+  }
+
+ private:
+  bool *hadFailure;
+
+  bool failableRewrite(FakeQuantOp op, PatternRewriter &rewriter) const {
+    auto converter =
+        mlir::quant::ir::ExpressedToQuantizedConverter::forInputType(
+            op.getType());
+    if (!converter) {
+      return (op.emitError("unsupported quantized type conversion"), true);
+    }
+
+    quant::QuantizedType elementType =
+        static_cast<const ConcreteRewriteClass *>(this)
+            ->convertFakeQuantAttrsToType(op, converter.expressed_type);
+
+    if (!elementType) {
+      // Note that the fakeQuantAttrsToType will have emitted the error.
+      return true;
+    }
+
+    Type quantizedType = converter.convert(elementType);
+    assert(quantizedType &&
+           "Converter accepted a type that it did not convert");
+
+    // TODO: Map to a qbarrier with an attribute like [Forced] to signal that
+    // this is a forced/hard-coded constraint.
+    auto qbarrier = QuantizeCastOp::create(rewriter, op.getLoc(), quantizedType,
+                                           op.getInputs());
+    rewriter.replaceOpWithNewOp<DequantizeCastOp>(op, converter.input_type,
+                                                  qbarrier.getResult());
+
+    return false;
+  }
+};
+
+class ConstFakeQuantRewrite
+    : public FakeQuantRewrite<ConstFakeQuantRewrite, ConstFakeQuant> {
+ public:
+  using BaseRewrite = FakeQuantRewrite<ConstFakeQuantRewrite, ConstFakeQuant>;
+
+  ConstFakeQuantRewrite(MLIRContext *ctx, bool *hadFailure)
+      : BaseRewrite(ctx, hadFailure) {}
+
+  quant::QuantizedType convertFakeQuantAttrsToType(ConstFakeQuant fqOp,
+                                                   Type expressedType) const {
+    return fakeQuantAttrsToType(
+        fqOp.getLoc(), fqOp.getNumBits(), fqOp.getMin().convertToFloat(),
+        fqOp.getMax().convertToFloat(), fqOp.getNarrowRange(), expressedType,
+        fqOp.getIsSigned());
+  }
+};
+
+class ConstFakeQuantPerAxisRewrite
+    : public FakeQuantRewrite<ConstFakeQuantPerAxisRewrite,
+                              ConstFakeQuantPerAxis> {
+ public:
+  using BaseRewrite =
+      FakeQuantRewrite<ConstFakeQuantPerAxisRewrite, ConstFakeQuantPerAxis>;
+
+  ConstFakeQuantPerAxisRewrite(MLIRContext *ctx, bool *hadFailure)
+      : BaseRewrite(ctx, hadFailure) {}
+
+  quant::QuantizedType convertFakeQuantAttrsToType(ConstFakeQuantPerAxis fqOp,
+                                                   Type expressedType) const {
+    SmallVector<double, 4> min, max;
+    min.reserve(fqOp.getMin().size());
+    max.reserve(fqOp.getMax().size());
+    for (auto m : fqOp.getMin())
+      min.push_back(cast<FloatAttr>(m).getValueAsDouble());
+    for (auto m : fqOp.getMax())
+      max.push_back(cast<FloatAttr>(m).getValueAsDouble());
+
+    return fakeQuantAttrsToType(fqOp.getLoc(), fqOp.getNumBits(),
+                                fqOp.getAxis(), min, max, fqOp.getNarrowRange(),
+                                expressedType, fqOp.getIsSigned());
+  }
+};
+
+void ConvertSimulatedQuantPass::runOnOperation() {
+  bool hadFailure = false;
+  auto func = getOperation();
+  RewritePatternSet patterns(func.getContext());
+  auto *ctx = func.getContext();
+  patterns.add<ConstFakeQuantRewrite, ConstFakeQuantPerAxisRewrite>(
+      ctx, &hadFailure);
+  (void)applyPatternsGreedily(func, std::move(patterns));
+  if (hadFailure) signalPassFailure();
+}
+
+std::unique_ptr<OperationPass<func::FuncOp>> createConvertSimulatedQuantPass() {
+  return std::make_unique<ConvertSimulatedQuantPass>();
+}
+
+}  // namespace mlir::quantfork

tensorflow/compiler/mlir/lite/quantization/ir/Passes.h

@@ -0,0 +1,57 @@
+/* Copyright 2022 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 file defines all of the passes owned by the quantization dialect. As
+// things mature, it is expected that passes specific to certain frontend or
+// backend dialects will move to those dialects directly. For now, they are
+// incubated here.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_IR_PASSES_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_IR_PASSES_H_
+
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+
+namespace mlir {
+namespace func {
+class FuncOp;
+}  // namespace func
+
+namespace quantfork {
+
+/// Creates a pass that converts quantization simulation operations (i.e.
+/// FakeQuant and those like it) to casts into/out of supported QuantizedTypes.
+std::unique_ptr<OperationPass<func::FuncOp>> createConvertSimulatedQuantPass();
+
+/// Creates a pass that converts constants followed by a qbarrier to a
+/// constant whose value is quantized. This is typically one of the last
+/// passes done when lowering to express actual quantized arithmetic in a
+/// low level representation. Because it modifies the constant, it is
+/// destructive and cannot be undone.
+std::unique_ptr<OperationPass<func::FuncOp>> createConvertConstPass();
+
+//===----------------------------------------------------------------------===//
+// Registration
+//===----------------------------------------------------------------------===//
+
+/// Generate the code for registering passes.
+#define GEN_PASS_REGISTRATION
+#include "tensorflow/compiler/mlir/lite/quantization/ir/Passes.h.inc"
+
+}  // namespace quantfork
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_IR_PASSES_H_

tensorflow/compiler/mlir/lite/quantization/ir/Passes.td

@@ -0,0 +1,34 @@
+/* Copyright 2022 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.
+==============================================================================*/
+
+#ifndef QUANT_FORK_PASSES
+#define QUANT_FORK_PASSES
+
+include "mlir/Pass/PassBase.td"
+
+def QuantConvertConst : Pass<"quant-convert-const", "func::FuncOp"> {
+  let summary = "Converts constants followed by qbarrier to actual quantized "
+                "values";
+  let constructor = "mlir::quantfork::createConvertConstPass()";
+}
+
+def QuantConvertSimulatedQuant
+    : Pass<"quant-convert-simulated-quantization", "func::FuncOp"> {
+  let summary = "Converts training-time simulated quantization ops to "
+                "corresponding quantize/dequantize casts";
+  let constructor = "mlir::quantfork::createConvertSimulatedQuantPass()";
+}
+
+#endif // QUANT_FORK_PASSES

tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.cc

@@ -0,0 +1,143 @@
+/* Copyright 2022 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 "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
+
+#include <cstdint>
+#include <functional>
+#include <iterator>
+#include <numeric>
+
+#include "llvm/ADT/STLExtras.h"
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+using namespace mlir;
+using namespace mlir::quantfork;
+
+using mlir::quant::QuantizedType;
+
+#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOpsDialect.cc.inc"
+
+void QuantizationForkDialect::initialize() {
+  addOperations<
+#define GET_OP_LIST
+#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.cc.inc"
+      >();
+}
+
+OpFoldResult StorageCastOp::fold(FoldAdaptor) {
+  // Matches x -> [scast -> scast] -> y, replacing the second scast with the
+  // value of x if the casts invert each other.
+  auto srcScastOp = getArg().getDefiningOp<StorageCastOp>();
+  if (!srcScastOp || srcScastOp.getArg().getType() != getType())
+    return OpFoldResult();
+  return srcScastOp.getArg();
+}
+
+/// The quantization specification should match the expressed type.
+static bool isValidQuantizationSpec(Attribute quantSpec, Type expressed) {
+  if (auto typeAttr = mlir::dyn_cast<TypeAttr>(quantSpec)) {
+    Type spec = typeAttr.getValue();
+    if (mlir::isa<TensorType, VectorType>(spec)) return false;
+
+    // The spec should be either a quantized type which is compatible to the
+    // expressed type, or a primitive type which is as same as the
+    // (element type of) the expressed type.
+    if (auto quantizedType = mlir::dyn_cast<QuantizedType>(spec))
+      return quantizedType.isCompatibleExpressedType(expressed);
+
+    if (auto tensorType = mlir::dyn_cast<TensorType>(expressed))
+      return spec == tensorType.getElementType();
+
+    if (auto vectorType = mlir::dyn_cast<VectorType>(expressed))
+      return spec == vectorType.getElementType();
+  }
+  return false;
+}
+
+LogicalResult QuantizeRegionOp::verify() {
+  // There are specifications for both inputs and outputs.
+  if (getNumOperands() != getInputSpecs().size() ||
+      getNumResults() != getOutputSpecs().size())
+    return emitOpError(
+        "has unmatched operands/results number and spec attributes number");
+
+  // Verify that quantization specifications are valid.
+  for (auto input : llvm::zip(getOperandTypes(), getInputSpecs())) {
+    Type inputType = std::get<0>(input);
+    Attribute inputSpec = std::get<1>(input);
+    if (!isValidQuantizationSpec(inputSpec, inputType)) {
+      return emitOpError() << "has incompatible specification " << inputSpec
+                           << " and input type " << inputType;
+    }
+  }
+
+  for (auto result : llvm::zip(getResultTypes(), getOutputSpecs())) {
+    Type outputType = std::get<0>(result);
+    Attribute outputSpec = std::get<1>(result);
+    if (!isValidQuantizationSpec(outputSpec, outputType)) {
+      return emitOpError() << "has incompatible specification " << outputSpec
+                           << " and output type " << outputType;
+    }
+  }
+  return success();
+}
+
+LogicalResult StatisticsOp::verify() {
+  auto tensorArg = mlir::dyn_cast<TensorType>(getArg().getType());
+  if (!tensorArg) return emitOpError("arg needs to be tensor type.");
+
+  // Verify layerStats attribute.
+  {
+    auto layerStatsType = getLayerStats().getShapedType();
+    if (!mlir::isa<FloatType>(layerStatsType.getElementType())) {
+      return emitOpError("layerStats must have a floating point element type");
+    }
+    if (layerStatsType.getRank() != 1 || layerStatsType.getDimSize(0) != 2) {
+      return emitOpError("layerStats must have shape [2]");
+    }
+  }
+  // Verify axisStats (optional) attribute.
+  if (getAxisStats()) {
+    if (!getAxis()) return emitOpError("axis must be specified for axisStats");
+
+    auto shape = tensorArg.getShape();
+    auto argSliceSize =
+        std::accumulate(std::next(shape.begin(), *getAxis()), shape.end(), 1,
+                        std::multiplies<int64_t>());
+
+    auto axisStatsType = getAxisStats()->getShapedType();
+    if (!mlir::isa<FloatType>(axisStatsType.getElementType())) {
+      return emitOpError("axisStats must have a floating point element type");
+    }
+    if (axisStatsType.getRank() != 2 || axisStatsType.getDimSize(1) != 2 ||
+        axisStatsType.getDimSize(0) != argSliceSize) {
+      return emitOpError(
+          "axisStats must have shape [N,2] "
+          "where N = the slice size defined by the axis dim");
+    }
+  }
+  return success();
+}
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.cc.inc"

tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h

@@ -0,0 +1,34 @@
+/* Copyright 2022 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_IR_QUANTOPS_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_IR_QUANTOPS_H_
+
+#include "llvm/Support/MathExtras.h"
+#include "mlir/Bytecode/BytecodeOpInterface.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Dialect.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Interfaces/InferTypeOpInterface.h"  // from @llvm-project
+#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOpsDialect.h.inc"
+#define GET_OP_CLASSES
+
+#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h.inc"
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_IR_QUANTOPS_H_

tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.td

@@ -0,0 +1,294 @@
+/* Copyright 2022 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 is the operation definition file for Quantization.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef QUANT_FORK_OPS
+#define QUANT_FORK_OPS
+
+include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOpsBase.td"
+include "mlir/Dialect/Quant/IR/QuantBase.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+
+class quant_TypedPrimitiveOrContainer<Type etype> :
+    Type<Or<[etype.predicate,
+                TensorOf<[etype]>.predicate,
+                VectorOfNonZeroRankOf<[etype]>.predicate]>,
+         "primitive/tensor/vector of " # etype.summary>;
+
+// A primitive type that can represent a real value. This is either a
+// floating point value or a quantized type.
+def quant_RealPrimitiveType :
+    Type<Or<[AnyFloat.predicate, quant_QuantizedType.predicate]>,
+    "real valued primitive (float or quantized type)">;
+
+// A primitive type that can represent a storage value. This is either an
+// integer or quantized type.
+def quant_StoragePrimitiveType :
+    Type<Or<[AnySignlessInteger.predicate, quant_QuantizedType.predicate]>,
+    "quantized storage primitive (integer or quantized type)">;
+
+// A primitive or container of RealPrimitiveType.
+def quant_RealValueType :
+    quant_TypedPrimitiveOrContainer<quant_RealPrimitiveType>;
+
+// A primitive or container of StoragePrimitiveType.
+def quant_StorageValueType :
+    quant_TypedPrimitiveOrContainer<quant_StoragePrimitiveType>;
+
+// Either a real valued or storage primitive or container type.
+def quant_RealOrStorageValueType :
+    Type<Or<[quant_RealValueType.predicate, quant_StorageValueType.predicate]>,
+    "real valued or storage primitive or container type">;
+
+//===----------------------------------------------------------------------===//
+// Base classes
+//===----------------------------------------------------------------------===//
+
+class quantfork_Op<string mnemonic, list<Trait> traits> :
+    Op<QuantizationFork_Dialect, mnemonic, traits>;
+
+//===----------------------------------------------------------------------===//
+// Quantization casts
+//===----------------------------------------------------------------------===//
+// A QuantizeCast (qcast) represents a potential type shift from a quantizable
+// type to a quantized type.
+//
+// At runtime, a qcast will apply the transformation expressed by its
+// operand and result type. For flexibility during transformation, it is also
+// possible to have a qcast that performs no transformation (both its
+// operand and result type are quantizable).
+//
+// A qcast will typically originate from either:
+//   a) An expressed or implied constraint in the source dialect which signals
+//      that a certain level of quantization is possible or required.
+//   b) An inference made by a quantization algorithm indicating that a
+//      quantized representation may be acceptable.
+//
+// Especially early in transformation, it is common to have pairs of
+// qcast/dcast at points where a transition to a quantized type is
+// required. In addition, it is also common to have an identity qcast
+// (where the operand and result type are not quantized) at all points where
+// it is legal to use a quantized representation (but is not known to be
+// acceptable).
+def quantfork_QuantizeCastOp : quantfork_Op<"qcast", [Pure]> {
+  let arguments = (ins quant_RealValueType:$arg);
+  let results = (outs quant_RealValueType);
+}
+
+// A DequantizeCast op (dcast) represents the inverse of a qcast,
+// converting back from a quantized to quantizable (expressed) type.
+//
+// Like qcasts, a dcast is allowed to have both its operand and result
+// as non quantized types. This facilitates transformations and marks edges
+// where the computation must be carried out in the expressed type.
+//
+// Especially early in transformation, it is common to have dcasts on
+// all operands to ops that must operate with the expressed type (typically
+// math ops prior to lowering to target-specific, quantized kernels).
+def quantfork_DequantizeCastOp : quantfork_Op<"dcast", [Pure]> {
+  let arguments = (ins quant_RealValueType:$arg);
+  let results = (outs quant_RealValueType);
+}
+
+// A StorageCast (scast) represents a cast from or to a type based on the
+// storage type and a type based on a corresponding quantized type.
+//
+// This op exists to ensure type coherency for between parts of the computation
+// which are operating directly on an underlying storage type and those which
+// operate on quantized values.
+//
+// Examples from storage to quantized type:
+//   i8 -> !quant<"uniform[i8:f32]{1.0}">
+//   tensor<4xi8> -> tensor<4x!quant<"uniform[i8:f32]{1.0}">>
+//   vector<4xi8> -> vector<4x!quant<"uniform[i8:f32]{1.0}">>
+def quantfork_StorageCastOp : quantfork_Op<"scast", [Pure]> {
+  let arguments = (ins quant_RealOrStorageValueType:$arg);
+  let results = (outs quant_RealOrStorageValueType);
+  let hasFolder = 1;
+}
+
+// A QuantizeRegion (region) represents a quantization unit which wraps
+// high-precision ops with quantization specifications for all the inputs
+// and outputs. Some quantization specifications can be undetermined and
+// derived from other ports by the target specification of the kernel.
+def quantfork_QuantizeRegionOp : quantfork_Op<"region", [
+    Pure,
+    IsolatedFromAbove,
+    SingleBlockImplicitTerminator<"ReturnOp">]> {
+  let summary = [{
+    The `region` operation wraps high-precision ops as a logical low-precision
+    quantized kernel.
+  }];
+
+  let arguments = (ins Variadic<AnyType>:$inputs,
+                    TypeArrayAttr:$input_specs,
+                    TypeArrayAttr:$output_specs,
+                    StrAttr:$logical_kernel);
+  let results = (outs Variadic<AnyType>:$outputs);
+  let regions = (region SizedRegion<1>:$body);
+  let hasVerifier = 1;
+}
+
+def quantfork_ReturnOp : quantfork_Op<"return", [Terminator]> {
+  let summary = [{
+    The `return` operation terminates a quantize region and returns values.
+  }];
+
+  let arguments = (ins Variadic<AnyTensor>:$results);
+}
+
+//===----------------------------------------------------------------------===//
+// Training integration and instrumentation ops
+//===----------------------------------------------------------------------===//
+
+def quantfork_ConstFakeQuant : quantfork_Op<"const_fake_quant",
+                                    [SameOperandsAndResultType, Pure]> {
+  let summary = [{
+    Simulates the effect of uniform quantization with const range.
+  }];
+
+  let description = [{
+    Given a const min, max, num_bits and narrow_range attribute, applies the
+    same uniform quantization simulation as is done by the TensorFlow
+    fake_quant_with_min_max_args op. See the fakeQuantAttrsToType() utility
+    method and the quant-convert-simulated-quantization pass for further details.
+  }];
+
+  let arguments = (ins
+    F32Tensor:$inputs,
+    F32Attr:$min,
+    F32Attr:$max,
+    // The bitwidth of the quantization; between 2 and 16, inclusive.
+    I64Attr:$num_bits,
+    // Quantization range starts from 0 or 1; starts from 1 if true.
+    DefaultValuedOptionalAttr<BoolAttr, "false">:$narrow_range,
+    // The sign of the quantization.
+    DefaultValuedOptionalAttr<BoolAttr, "false">:$is_signed
+  );
+
+  let results = (outs
+    F32Tensor:$outputs
+  );
+}
+
+def quantfork_ConstFakeQuantPerAxis : quantfork_Op<"const_fake_quant_per_axis",
+                                    [SameOperandsAndResultType, Pure]> {
+  let summary = [{
+    Simulates the effect of per axis uniform quantization with const range.
+  }];
+
+  let description = [{
+    Given a const min, max, num_bits and narrow_range attribute, applies the
+    same per axis uniform quantization simulation as is done by the TensorFlow
+    fake_quant_with_min_max_vars_per_channel op. See the fakeQuantAttrsToType()
+    utility method and the quant-convert-simulated-quantization pass for further
+    details.
+  }];
+
+  let arguments = (ins
+    F32Tensor:$inputs,
+    F32ArrayAttr:$min,
+    F32ArrayAttr:$max,
+    // The quantized dimension of the inputs tensor.
+    I64Attr:$axis,
+    // The bitwidth of the quantization; between 2 and 16, inclusive.
+    I64Attr:$num_bits,
+    // Quantization range starts from 0 or 1; starts from 1 if true.
+    DefaultValuedOptionalAttr<BoolAttr, "false">:$narrow_range,
+    // The sign of the quantization.
+    DefaultValuedOptionalAttr<BoolAttr, "false">:$is_signed
+  );
+
+  let results = (outs
+    F32Tensor:$outputs
+  );
+}
+
+def quantfork_StatisticsRefOp : quantfork_Op<"stats_ref", [SameOperandsAndResultType]> {
+  let summary = "Indicates that statistics are resolved by reference.";
+
+  let description = [{
+    This op acts as an identity that, when encountered at runtime, should result
+    in statistics being collected about about the value of its operand/result.
+    Such statistics will be stored with the provided key, allowing this node
+    to later be converted to a 'stats' op if statistics with that key have been
+    encountered.
+  }];
+
+  let arguments = (ins
+    quant_RealValueType:$arg,
+    StrAttr:$statsKey
+  );
+  let results = (outs quant_RealValueType);
+}
+
+def quantfork_StatisticsOp : quantfork_Op<"stats", [SameOperandsAndResultType]> {
+  let summary = "Identity op which associates statistics with the value.";
+
+  let description = [{
+    Associates statistics about the runtime ranges of values observed for
+    evaluations of this node.
+
+    Statistics about the entire type are reported in the 'layerStats' attribute
+    and those for each axis, in the (optional) `axisStats` attribute. The
+    interpretation of each is determined by the last dimension of its shape.
+    Currently, only dim=2 is supported, which is interpreted as [min, max].
+
+    `layerStats` must be a rank 1 tensor: [2]
+    `axisStats` must be a rank 2 tensor: [N, 2], where N=the slice size
+      splitted by the `axis` dimension. For example:
+
+    ```
+    <?x?x3x2>, axis=3 => N=2
+    <?x?x3x2>, axis=2 => N=6
+    ```
+  }];
+
+  let arguments = (ins
+    quant_RealValueType:$arg,
+    ElementsAttr:$layerStats,
+    OptionalAttr<ElementsAttr>:$axisStats,
+    OptionalAttr<I64Attr>:$axis);
+  let results = (outs quant_RealValueType);
+  let hasVerifier = 1;
+}
+
+def quantfork_CoupledRefOp : quantfork_Op<"coupled_ref", [SameOperandsAndResultType]> {
+  let summary = [{
+    Indicates that one point of the computation is coupled to another.
+  }];
+
+  let description = [{
+    Ordinarily, relationships between ops for the purposes of determining
+    compatible quantized types is explicit based on the use-def chain. However,
+    in some situations, a use may be separated from its def by arbitrary
+    external connections. In such a case, during analysis, all coupled_ref
+    nodes in a module which share a coupledKey will be considered to be
+    directly connected as via an identity op for the purpose of type inference.
+  }];
+
+  let arguments = (ins
+    quant_RealValueType:$arg,
+    StrAttr:$coupledKey);
+  let results = (outs quant_RealValueType);
+}
+
+#endif // QUANT_FORK_OPS

tensorflow/compiler/mlir/lite/quantization/ir/QuantOpsBase.td

@@ -0,0 +1,32 @@
+/* Copyright 2022 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.
+==============================================================================*/
+//
+//===----------------------------------------------------------------------===//
+//
+// Predicates for types in the Quantization dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef QUANT_FORK_BASE
+#define QUANT_FORK_BASE
+
+include "mlir/IR/OpBase.td"
+
+def QuantizationFork_Dialect : Dialect {
+  let name = "quantfork";
+  let cppNamespace = "::mlir::quantfork";
+}
+
+#endif // QUANT_FORK_BASE

tensorflow/compiler/mlir/lite/quantization/ir/QuantizeUtils.cc

@@ -0,0 +1,147 @@
+/* Copyright 2022 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 "tensorflow/compiler/mlir/lite/quantization/ir/QuantizeUtils.h"
+
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.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;
+
+/// Converts a possible primitive, real expressed value attribute to a
+/// corresponding storage attribute (typically FloatAttr -> IntegerAttr).
+/// quantizedElementType is the QuantizedType that describes the expressed
+/// origValue.
+/// 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) {
+  if (mlir::isa<FloatAttr>(origRealValue)) {
+    FloatAttr floatAttr = mlir::cast<FloatAttr>(origRealValue);
+    outConvertedType = quantizedElementType.getStorageType();
+    return IntegerAttr::get(quantizedElementType.getStorageType(),
+                            converter.quantizeFloatToInt(floatAttr.getValue()));
+  }
+
+  return nullptr;
+}
+
+/// Converts a real expressed DenseFPElementsAttr to a corresponding
+/// DenseElementsAttr (typically DenseIntElementsAttr) containing quantized
+/// storage values assuming the given quantizedElementType and converter.
+static DenseElementsAttr convertDenseFPElementsAttr(
+    DenseFPElementsAttr realFPElementsAttr,
+    quant::QuantizedType quantizedElementType,
+    const mlir::quant::ir::UniformQuantizedValueConverter &converter) {
+  return realFPElementsAttr.mapValues(
+      quantizedElementType.getStorageType(),
+      [&converter](const APFloat &realVal) {
+        return converter.quantizeFloatToInt(realVal);
+      });
+}
+
+/// Converts a real expressed SplatElementsAttr to a corresponding
+/// SplatElementsAttr containing quantized storage values assuming the given
+/// quantizedElementType and converter.
+static SparseElementsAttr convertSparseElementsAttr(
+    SparseElementsAttr realSparseAttr,
+    quant::QuantizedType quantizedElementType,
+    const mlir::quant::ir::UniformQuantizedValueConverter &converter) {
+  DenseElementsAttr realDenseAttr = realSparseAttr.getValues();
+  if (!mlir::isa<DenseFPElementsAttr>(realDenseAttr)) {
+    return nullptr;
+  }
+  DenseElementsAttr quantDenseAttr =
+      convertDenseFPElementsAttr(mlir::cast<DenseFPElementsAttr>(realDenseAttr),
+                                 quantizedElementType, converter);
+  if (!quantDenseAttr) {
+    return nullptr;
+  }
+
+  // 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<ShapedType>(
+      quantizedElementType.castExpressedToStorageType(
+          realSparseAttr.getType()));
+  if (!newSparseType) {
+    return nullptr;
+  }
+  return SparseElementsAttr::get(newSparseType, realSparseAttr.getIndices(),
+                                 quantDenseAttr);
+}
+
+/// 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 realValue, quant::UniformQuantizedType quantizedElementType,
+    const mlir::quant::ir::UniformQuantizedValueConverter &converter,
+    Type &outConvertedType) {
+  // Fork to handle different variants of constants supported.
+  if (mlir::isa<DenseFPElementsAttr>(realValue)) {
+    // Dense tensor or vector constant.
+    auto converted =
+        convertDenseFPElementsAttr(mlir::cast<DenseFPElementsAttr>(realValue),
+                                   quantizedElementType, converter);
+    outConvertedType = converted.getType();
+    return converted;
+  }
+  if (mlir::isa<SparseElementsAttr>(realValue)) {
+    // Sparse tensor or vector constant.
+    auto converted =
+        convertSparseElementsAttr(mlir::cast<SparseElementsAttr>(realValue),
+                                  quantizedElementType, converter);
+    outConvertedType = converted.getType();
+    return converted;
+  }
+  // Nothing else matched: try to convert a primitive.
+  return convertPrimitiveValueAttr(realValue, quantizedElementType, converter,
+                                   outConvertedType);
+}
+
+/// Convert an attribute from a type based on
+/// quantizedElementType.getExpressedType() to one based on
+/// 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 =
+          mlir::dyn_cast<quant::UniformQuantizedType>(quantizedElementType)) {
+    mlir::quant::ir::UniformQuantizedValueConverter converter(uniformQuantized);
+    return quantizeAttrUniform(realValue, uniformQuantized, converter,
+                               outConvertedType);
+  }
+  if (auto uniformQuantizedPerAxis =
+          mlir::dyn_cast<quant::UniformQuantizedPerAxisType>(
+              quantizedElementType)) {
+    mlir::quant::ir::UniformQuantizedPerAxisValueConverter converter(
+        uniformQuantizedPerAxis);
+    auto converted = converter.convert(realValue);
+    // TODO: why we need this outConvertedType? remove it?
+    if (converted) {
+      outConvertedType = converted.getType();
+    }
+    return converted;
+  }
+  return nullptr;
+}

tensorflow/compiler/mlir/lite/quantization/ir/QuantizeUtils.h

@@ -0,0 +1,73 @@
+/* Copyright 2022 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_IR_QUANTIZEUTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_IR_QUANTIZEUTILS_H_
+
+namespace mlir {
+class Attribute;
+class Type;
+
+namespace quant {
+class QuantizedType;
+class UniformQuantizedType;
+namespace ir {
+class UniformQuantizedValueConverter;
+}  // namespace ir
+}  // namespace quant
+namespace quantfork {
+
+/// Converts an attribute from a type based on
+/// quantizedElementType.getExpressedType() to one based on
+/// quantizedElementType.getStorageType(), where quantizedElementType is as from
+/// QuantizedType::getQuantizedElementType().
+/// Returns nullptr if the conversion is not supported. On success, stores the
+/// converted type in outConvertedType.
+///
+/// Examples:
+/// 1. realValue is a primitive value attribute:
+/// (realValue: FloatAttr, quantizedElementType: UniformQuantizedType[i8:f32])
+///   -> (IntegerAttr, outConvertedType: i8)
+/// 2. realValue is an elements attribute:
+/// (realValue: DenseElementsAttr[tensor<2x2xf32>],
+///  quantizedElementType: UniformQuantizedType[i8:f32])
+///   -> (DenseElementsAttr[tensor<2x2xi8>], outConvertedType: tensor<2x2xi8>)
+Attribute quantizeAttr(Attribute realValue,
+                       quant::QuantizedType quantizedElementType,
+                       Type &outConvertedType);
+
+/// Converts an attribute from a type based on
+/// quantizedElementType.getExpressedType() to one based on
+/// quantizedElementType.getStorageType(), where quantizedElementType is as from
+/// QuantizedType::getQuantizedElementType() and casted to an
+/// UniformQuantizedType. Returns nullptr if the conversion is not supported. On
+/// success, stores the converted type in outConvertedType.
+///
+/// Examples:
+/// 1. realValue is a primitive value attribute:
+/// (realValue: FloatAttr, quantizedElementType: UniformQuantizedType[i8:f32])
+///   -> (IntegerAttr, outConvertedType: i8)
+/// 2. realValue is an elements attribute:
+/// (realValue: DenseElementsAttr[tensor<2x2xf32>],
+///  quantizedElementType: UniformQuantizedType[i8:f32])
+///   -> (DenseElementsAttr[tensor<2x2xi8>], outConvertedType: tensor<2x2xi8>)
+Attribute quantizeAttrUniform(
+    Attribute realValue, quant::UniformQuantizedType quantizedElementType,
+    const mlir::quant::ir::UniformQuantizedValueConverter &converter,
+    Type &outConvertedType);
+}  // namespace quantfork
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_QUANTIZATION_IR_QUANTIZEUTILS_H_