chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/quantization/common/ir/BUILD

@@ -0,0 +1,91 @@
+load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+
+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=quantization",
+        ],
+        "QuantOpsDialect.cc.inc": [
+            "-gen-dialect-defs",
+            "-dialect=quantization",
+        ],
+    },
+    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=tfquant",
+    ]},
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "Passes.td",
+    deps = ["@llvm-project//mlir:PassBaseTdFiles"],
+)
+
+cc_library(
+    name = "QuantOps",
+    srcs = [
+        "ConvertConst.cc",
+        "ConvertSimQuant.cc",
+        "FakeQuantSupport.cc",
+        "QuantOps.cc",
+        "QuantizeUtils.cc",
+        "UniformSupport.cc",
+    ],
+    hdrs = [
+        "FakeQuantSupport.h",
+        "Passes.h",
+        "QuantOps.h",
+        "QuantizeUtils.h",
+        "UniformSupport.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        ":QuantOpsIncGen",
+        ":QuantPassIncGen",
+        "@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/quantization/common/ir/ConvertConst.cc

@@ -0,0 +1,124 @@
+/* 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/quantization/common/ir/Passes.h"
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantizeUtils.h"
+
+namespace mlir {
+namespace quant::ir {
+
+using mlir::quant::QuantizedType;
+
+namespace {
+#define GEN_PASS_DEF_QUANTCONVERTCONST
+#include "tensorflow/compiler/mlir/quantization/common/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>> createConvertConstPass() {
+  return std::make_unique<ConvertConstPass>();
+}
+
+}  // namespace quant::ir
+}  // namespace mlir

tensorflow/compiler/mlir/quantization/common/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/quantization/common/ir/FakeQuantSupport.h"
+#include "tensorflow/compiler/mlir/quantization/common/ir/Passes.h"
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/quantization/common/ir/UniformSupport.h"
+
+namespace mlir::quant::ir {
+
+#define GEN_PASS_DEF_QUANTCONVERTSIMULATEDQUANT
+#include "tensorflow/compiler/mlir/quantization/common/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 quantfork::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 quantfork::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::quant::ir

tensorflow/compiler/mlir/quantization/common/ir/FakeQuantSupport.cc

@@ -0,0 +1,214 @@
+/* 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/quantization/common/ir/FakeQuantSupport.h"
+
+#include <cassert>
+#include <cmath>
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Diagnostics.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir::quantfork {
+
+static bool getDefaultStorageParams(unsigned numBits, bool narrowRange,
+                                    bool isSigned, MLIRContext *ctx,
+                                    Type &storageType, int64_t &qmin,
+                                    int64_t &qmax) {
+  // Hard-coded type mapping from TFLite.
+  if (numBits <= 4) {
+    storageType = IntegerType::get(ctx, 4);
+    if (isSigned) {
+      qmin = -8;
+      qmax = 7;
+    } else {
+      qmin = 0;
+      qmax = 15;
+    }
+  } else if (numBits <= 8) {
+    storageType = IntegerType::get(ctx, 8);
+    if (isSigned) {
+      qmin = -128;
+      qmax = 127;
+    } else {
+      qmin = 0;
+      qmax = 255;
+    }
+  } else if (numBits <= 16) {
+    storageType = IntegerType::get(ctx, 16);
+    if (isSigned) {
+      qmin = -32768;
+      qmax = 32767;
+    } else {
+      qmin = 0;
+      qmax = 65535;
+    }
+  } else if (numBits <= 32) {
+    storageType = IntegerType::get(ctx, 32);
+    if (isSigned) {
+      qmin = std::numeric_limits<int32_t>::min();
+      qmax = std::numeric_limits<int32_t>::max();
+    } else {
+      qmin = std::numeric_limits<uint32_t>::min();
+      qmax = std::numeric_limits<uint32_t>::max();
+    }
+  } else {
+    return true;
+  }
+
+  // Handle narrowRange.
+  if (narrowRange) {
+    qmin += 1;
+  }
+  return false;
+}
+
+// This is a specific implementation of nudging:
+// If 0.0 < rmin < rmax or rmin < rmax < 0.0, the range will be shifted
+// to include 0.0, but the range width size (rmax-rmin) isn't changed. The zero
+// point is derived from the shifted range, and the scale isn't changed. As
+// a consequence some values, which are supposed in the original [rmin, rmax]
+// range will be outside the shifted range and be clamped during quantization.
+// TODO: we should nudge the scale as well, but that requires the
+// fake quant op used in the training to use the nudged scale as well.
+static void getNudgedScaleAndZeroPoint(int64_t qmin, int64_t qmax, double rmin,
+                                       double rmax, double &scale,
+                                       int64_t &nudgedZeroPoint) {
+  // Determine the scale.
+  const double qminDouble = qmin;
+  const double qmaxDouble = qmax;
+  scale = (rmax - rmin) / (qmaxDouble - qminDouble);
+
+  // Zero point computation.
+  // In float, solve the affine equation for any known pair
+  // (real value, corresponding quantized value), of which, two such pairs
+  // are known: (rmin, qmin), (rmax, qmax).
+  // The arithmetic error on the zero point computed from either pair will be
+  // roughly machine_epsilon * (sum of absolute values of terms).
+  // Use the variant that adds the smaller error.
+  const double zeroPointFromMin = qminDouble - rmin / scale;
+  const double zeroPointFromMinError =
+      std::abs(qminDouble) + std::abs(rmin / scale);
+  const double zeroPointFromMax = qmaxDouble - rmax / scale;
+  const double zeroPointFromMaxError =
+      std::abs(qmaxDouble) + std::abs(rmax / scale);
+
+  const double zeroPointDouble = (zeroPointFromMinError < zeroPointFromMaxError)
+                                     ? zeroPointFromMin
+                                     : zeroPointFromMax;
+
+  // Now nudge the zero point to be an integer.
+  nudgedZeroPoint = 0;
+  if (zeroPointDouble < qminDouble) {
+    nudgedZeroPoint = qmin;
+  } else if (zeroPointDouble > qmaxDouble) {
+    nudgedZeroPoint = qmax;
+  } else {
+    nudgedZeroPoint = round(zeroPointDouble);
+  }
+
+  // By construction, the nudged zero point should always be in range.
+  assert(nudgedZeroPoint >= qmin);
+  assert(nudgedZeroPoint <= qmax);
+}
+
+quant::UniformQuantizedType fakeQuantAttrsToType(Location loc, unsigned numBits,
+                                                 double rmin, double rmax,
+                                                 bool narrowRange,
+                                                 Type expressedType,
+                                                 bool isSigned) {
+  MLIRContext *ctx = expressedType.getContext();
+  unsigned flags = isSigned ? quant::QuantizationFlags::Signed : 0;
+  Type storageType;
+  int64_t qmin;
+  int64_t qmax;
+  if (getDefaultStorageParams(numBits, narrowRange, isSigned, ctx, storageType,
+                              qmin, qmax)) {
+    return (emitError(loc, "unsupported FakeQuant number of bits: ") << numBits,
+            nullptr);
+  }
+
+  // Special case where min/max is close enough. The tensor contents are all
+  // 0.0s, so the scale is set to 1.0 and the tensor can be quantized to zero
+  // points and dequantized to 0.0.
+  if (std::fabs(rmax - rmin) < std::numeric_limits<double>::epsilon()) {
+    return quant::UniformQuantizedType::getChecked(
+        loc, flags, storageType, expressedType, 1.0, qmin, qmin, qmax);
+  }
+
+  double scale;
+  int64_t nudgedZeroPoint;
+  getNudgedScaleAndZeroPoint(qmin, qmax, rmin, rmax, scale, nudgedZeroPoint);
+
+  return quant::UniformQuantizedType::getChecked(loc, flags, storageType,
+                                                 expressedType, scale,
+                                                 nudgedZeroPoint, qmin, qmax);
+}
+
+quant::UniformQuantizedPerAxisType fakeQuantAttrsToType(
+    Location loc, unsigned numBits, int32_t quantizedDimension,
+    ArrayRef<double> rmins, ArrayRef<double> rmaxs, bool narrowRange,
+    Type expressedType, bool isSigned) {
+  size_t axisSize = rmins.size();
+  if (axisSize != rmaxs.size()) {
+    return (emitError(loc, "mismatched per-axis min and max size: ")
+                << axisSize << " vs. " << rmaxs.size(),
+            nullptr);
+  }
+
+  MLIRContext *ctx = expressedType.getContext();
+  Type storageType;
+  int64_t qmin;
+  int64_t qmax;
+  if (getDefaultStorageParams(numBits, narrowRange, isSigned, ctx, storageType,
+                              qmin, qmax)) {
+    return (emitError(loc, "unsupported FakeQuant number of bits: ") << numBits,
+            nullptr);
+  }
+
+  SmallVector<double, 4> scales;
+  SmallVector<int64_t, 4> zeroPoints;
+  scales.reserve(axisSize);
+  zeroPoints.reserve(axisSize);
+  for (size_t axis = 0; axis != axisSize; ++axis) {
+    double rmin = rmins[axis];
+    double rmax = rmaxs[axis];
+    if (std::fabs(rmax - rmin) < std::numeric_limits<double>::epsilon()) {
+      scales.push_back(1.0);
+      zeroPoints.push_back(qmin);
+      continue;
+    }
+
+    double scale;
+    int64_t nudgedZeroPoint;
+    getNudgedScaleAndZeroPoint(qmin, qmax, rmin, rmax, scale, nudgedZeroPoint);
+    scales.push_back(scale);
+    zeroPoints.push_back(nudgedZeroPoint);
+  }
+
+  unsigned flags = isSigned ? quant::QuantizationFlags::Signed : 0;
+  return quant::UniformQuantizedPerAxisType::getChecked(
+      loc, flags, storageType, expressedType, scales, zeroPoints,
+      quantizedDimension, qmin, qmax);
+}
+
+}  // namespace mlir::quantfork

tensorflow/compiler/mlir/quantization/common/ir/FakeQuantSupport.h

@@ -0,0 +1,79 @@
+/* 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 support utilities for interoperating with FakeQuant* based
+// QAT (Quantized Aware Training) computations, as implemented by TFLite. Note
+// that FakeQuant* operators mix multiple concerns specific to how TFLite
+// originally implemented quantization. As such, utilities here enforce
+// opinions taken by that codebase (vs providing any amount of genericity).
+//
+// Specifically, it combines the following concerns, each of which would be
+// independent variables in a more generic setup:
+//   - numBits and isSigned imply storage data type (uint8, int8, int16)
+//   - numBits < 8 is promoted to uint8 or int8
+//   - "narrow_range" narrows the lower bound of the storage type's range by
+//     1
+//   - the specified min/max values are "nudged" so that the result has a zero
+//     that can be exactly expressed
+//   - min=max=0 implies scale=0 and zero_point=0
+//
+// With the above assumptions applied, every conforming specified FakeQuant op
+// can be represented by a UniformQuantizedType. This scheme is not expected to
+// be generalized further in the future and should be considered to be a
+// legacy set of rules.
+//
+// As canonically used in TensorFlow graphs, the presence of a FakeQuant node
+// is a hint that the specific math represented here has been simulated at
+// training time. As such, it is usually not advised to arbitrarily change
+// quantization parameters derived from FakeQuant.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_IR_FAKEQUANTSUPPORT_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_IR_FAKEQUANTSUPPORT_H_
+
+#include <cstdint>
+
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir {
+namespace quantfork {
+
+/// Converts per-layer FakeQuant attributes to the corresponding type.
+/// In the event that the parameters cannot be converted, returns a nullptr
+/// convertible Type and issues an appropriate error.
+/// Note that there are multiple variants of a per-layer FakeQuant op, so
+/// this function takes the attributes discretely vs taking a reference to the
+/// originating op.
+quant::UniformQuantizedType fakeQuantAttrsToType(Location loc, unsigned numBits,
+                                                 double rmin, double rmax,
+                                                 bool narrowRange,
+                                                 Type expressedType,
+                                                 bool isSigned = false);
+
+/// Converts per-channel FakeQuant attributes to the corresponding type.
+/// In the event that the parameters cannot be converted, returns a nullptr
+/// convertible Type and issues an appropriate error.
+quant::UniformQuantizedPerAxisType fakeQuantAttrsToType(
+    Location loc, unsigned numBits, int32_t quantizedDimension,
+    ArrayRef<double> rmins, ArrayRef<double> rmax, bool narrowRange,
+    Type expressedType, bool isSigned = false);
+}  // namespace quantfork
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_IR_FAKEQUANTSUPPORT_H_

tensorflow/compiler/mlir/quantization/common/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_QUANTIZATION_COMMON_IR_PASSES_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_IR_PASSES_H_
+
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+
+namespace mlir {
+namespace func {
+class FuncOp;
+}  // namespace func
+
+namespace quant::ir {
+
+/// 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/quantization/common/ir/Passes.h.inc"
+
+}  // namespace quant::ir
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_IR_PASSES_H_

tensorflow/compiler/mlir/quantization/common/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 TF_QUANT_PASSES
+#define TF_QUANT_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::quant::ir::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::quant::ir::createConvertSimulatedQuantPass()";
+}
+
+#endif // TF_QUANT_PASSES

tensorflow/compiler/mlir/quantization/common/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/quantization/common/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
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOpsDialect.cc.inc"
+
+namespace mlir::quant::ir {
+
+using mlir::quant::QuantizedType;
+
+void TFQuantDialect::initialize() {
+  addOperations<
+#define GET_OP_LIST
+#include "tensorflow/compiler/mlir/quantization/common/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();
+}
+
+}  // namespace mlir::quant::ir
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.cc.inc"

tensorflow/compiler/mlir/quantization/common/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_QUANTIZATION_COMMON_IR_QUANTOPS_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_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/quantization/common/ir/QuantOpsDialect.h.inc"
+#define GET_OP_CLASSES
+
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h.inc"
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_IR_QUANTOPS_H_

tensorflow/compiler/mlir/quantization/common/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 QUANTIZATION_OPS
+#define QUANTIZATION_OPS
+
+include "mlir/Dialect/Quant/IR/QuantBase.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "tensorflow/compiler/mlir/quantization/common/ir/QuantOpsBase.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 Quantization_Op<string mnemonic, list<Trait> traits> :
+    Op<TF_Quant_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 Quantization_QuantizeCastOp : Quantization_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 Quantization_DequantizeCastOp : Quantization_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 Quantization_StorageCastOp : Quantization_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 Quantization_QuantizeRegionOp : Quantization_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 Quantization_ReturnOp : Quantization_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 Quantization_ConstFakeQuant : Quantization_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 Quantization_ConstFakeQuantPerAxis : Quantization_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 Quantization_StatisticsRefOp : Quantization_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 Quantization_StatisticsOp : Quantization_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 Quantization_CoupledRefOp : Quantization_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 // Quantization_OPS

tensorflow/compiler/mlir/quantization/common/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 QUANTIZATION_BASE
+#define QUANTIZATION_BASE
+
+include "mlir/IR/OpBase.td"
+
+def TF_Quant_Dialect : Dialect {
+  let name = "quantization";
+  let cppNamespace = "::mlir::quant::ir";
+}
+
+#endif // QUANTIZATION_BASE

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

@@ -0,0 +1,148 @@
+/* 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/quantization/common/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"
+
+namespace mlir {
+namespace quant::ir {
+
+/// 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 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 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 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 quantizeAttrUniform(Attribute realValue,
+                              quant::UniformQuantizedType quantizedElementType,
+                              const 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 quantizeAttr(Attribute realValue,
+                       quant::QuantizedType quantizedElementType,
+                       Type &outConvertedType) {
+  if (auto uniformQuantized =
+          mlir::dyn_cast<quant::UniformQuantizedType>(quantizedElementType)) {
+    UniformQuantizedValueConverter converter(uniformQuantized);
+    return quantizeAttrUniform(realValue, uniformQuantized, converter,
+                               outConvertedType);
+  }
+  if (auto uniformQuantizedPerAxis =
+          mlir::dyn_cast<quant::UniformQuantizedPerAxisType>(
+              quantizedElementType)) {
+    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;
+}
+
+}  // namespace quant::ir
+}  // namespace mlir

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

@@ -0,0 +1,71 @@
+/* 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_QUANTIZATION_COMMON_IR_QUANTIZEUTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_IR_QUANTIZEUTILS_H_
+
+namespace mlir {
+class Attribute;
+class Type;
+
+namespace quant {
+class QuantizedType;
+
+namespace ir {
+class UniformQuantizedType;
+class UniformQuantizedValueConverter;
+
+/// 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, 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,
+                              UniformQuantizedType quantizedElementType,
+                              const UniformQuantizedValueConverter &converter,
+                              Type &outConvertedType);
+}  // namespace ir
+}  // namespace quant
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_IR_QUANTIZEUTILS_H_

tensorflow/compiler/mlir/quantization/common/ir/UniformSupport.cc

@@ -0,0 +1,112 @@
+/* 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/quantization/common/ir/UniformSupport.h"
+
+#include <cassert>
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+#include <iterator>
+#include <numeric>
+
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/Attributes.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/Types.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir::quant::ir {
+
+static bool isQuantizablePrimitiveType(Type input_type) {
+  return isa<FloatType>(input_type);
+}
+
+ExpressedToQuantizedConverter ExpressedToQuantizedConverter::forInputType(
+    Type input_type) {
+  if (isa<TensorType, VectorType>(input_type)) {
+    Type element_type = cast<ShapedType>(input_type).getElementType();
+    if (!isQuantizablePrimitiveType(element_type))
+      return ExpressedToQuantizedConverter{input_type, nullptr};
+    return ExpressedToQuantizedConverter{input_type, element_type};
+  }
+  // Supported primitive type (which just is the expressed type).
+  if (isQuantizablePrimitiveType(input_type))
+    return ExpressedToQuantizedConverter{input_type, input_type};
+  // Unsupported.
+  return ExpressedToQuantizedConverter{input_type, nullptr};
+}
+
+Type ExpressedToQuantizedConverter::convert(
+    quant::QuantizedType elemental_type) const {
+  assert(expressed_type && "convert() on unsupported conversion");
+  if (auto tensor_type = dyn_cast<RankedTensorType>(input_type))
+    return RankedTensorType::get(tensor_type.getShape(), elemental_type);
+  if (auto tensor_type = dyn_cast<UnrankedTensorType>(input_type))
+    return UnrankedTensorType::get(elemental_type);
+  if (auto vector_type = dyn_cast<VectorType>(input_type))
+    return VectorType::get(vector_type.getShape(), elemental_type);
+
+  // If the expressed types match, just use the new elemental type.
+  if (elemental_type.getExpressedType() == expressed_type) {
+    return elemental_type;
+  }
+  // Unsupported.
+  return nullptr;
+}
+
+ElementsAttr UniformQuantizedPerAxisValueConverter::convert(
+    Attribute real_value) {
+  if (auto attr = dyn_cast<DenseFPElementsAttr>(real_value)) {
+    return convert(attr);
+  }
+  return nullptr;
+}
+
+DenseElementsAttr UniformQuantizedPerAxisValueConverter::convert(
+    DenseFPElementsAttr attr) {
+  // Creates the converter for each chunk. Normally the size of the
+  // quantization dim is 3, so we can cache all the converters.
+  ShapedType type = attr.getType();
+  std::size_t dim_size = type.getDimSize(quantization_dim_);
+  if (dim_size != scales_.size()) {
+    return {};
+  }
+  SmallVector<UniformQuantizedValueConverter, 4> converters;
+  converters.reserve(dim_size);
+  for (int i = 0, e = dim_size; i != e; ++i) {
+    converters.push_back(getPerChunkConverter(i));
+  }
+
+  // Scan the elements of the dense elements attributes and quantize them by
+  // using the right quantization parameters.
+  int64_t flatten_index = 0;
+  auto shape = type.getShape();
+  int64_t chunk_size =
+      std::accumulate(std::next(shape.begin(), quantization_dim_ + 1),
+                      shape.end(), 1, std::multiplies<int64_t>());
+  Type new_element_type =
+      IntegerType::get(attr.getContext(), storage_bit_width_);
+  return attr.mapValues(new_element_type, [&](const APFloat &old) {
+    int chunk_index = flatten_index / chunk_size;
+    flatten_index++;
+    return converters[chunk_index % dim_size].quantizeFloatToInt(old);
+  });
+}
+
+}  // namespace mlir::quant::ir

tensorflow/compiler/mlir/quantization/common/ir/UniformSupport.h

@@ -0,0 +1,247 @@
+/* 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_QUANTIZATION_COMMON_IR_UNIFORMSUPPORT_H_
+#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_IR_UNIFORMSUPPORT_H_
+
+#include <algorithm>
+#include <cassert>
+#include <cmath>
+#include <cstdint>
+#include <cstring>
+#include <utility>
+
+#include "llvm/ADT/APFloat.h"
+#include "llvm/ADT/APInt.h"
+#include "llvm/ADT/APSInt.h"
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir::quant::ir {
+
+// Performs type conversion from an arbitrary input type to a type
+// that is expressed by a QuantizedType.
+//
+// This handles cases where the inputType is a supported primitive type
+// (i.e. f32, bf16, etc) or a vector/tensor type based on a supported
+// elemental type.
+//
+// Since conversion often involves introspecting some attributes of the
+// input type in order to determine how to represent it, this is a two step
+// process.
+struct ExpressedToQuantizedConverter {
+  // Creates a converter for the given input type.
+  static ExpressedToQuantizedConverter forInputType(Type input_type);
+
+  // Converts the inputType to be based on the given elemental type,
+  // returning the new type (or nullptr and emit an error on failure).
+  Type convert(quant::QuantizedType elemental_type) const;
+
+  // Whether the conversion is legal.
+  explicit operator bool() const { return (bool)expressed_type; }
+
+  // The input type that is being converted from.
+  // This may be an elemental or composite type.
+  const Type input_type;
+
+  // Supported, elemental expressed type (i.e. f32).
+  // Will be nullptr if conversion is not supported.
+  const Type expressed_type;
+};
+
+// Reference implementation of converting between real numbers and values
+// represented by a UniformQuantizedType.
+// Note that this is not expected to be speedy and may be superseded eventually
+// by a more optimal implementation.
+// Also, the interface assumes that quantization is done per-layer and will
+// need to be wider for various per-channel schemes. As such, this is a
+// placeholder.
+class UniformQuantizedValueConverter {
+ public:
+  explicit UniformQuantizedValueConverter(
+      quant::UniformQuantizedType uniform_type)
+      : UniformQuantizedValueConverter(
+            uniform_type.getScale(),
+            static_cast<double>(uniform_type.getZeroPoint()),
+            static_cast<double>(uniform_type.getStorageTypeMin()),
+            static_cast<double>(uniform_type.getStorageTypeMax()),
+            uniform_type.getStorageTypeIntegralWidth(),
+            uniform_type.isSigned()) {
+    assert(isa<FloatType>(uniform_type.getExpressedType()));
+    assert(uniform_type.getStorageType().isSignlessInteger());
+  }
+
+  UniformQuantizedValueConverter(double scale, double zero_point,
+                                 double clamp_min, double clamp_max,
+                                 uint32_t storage_bit_width, bool is_signed)
+      : scale_(scale),
+        zero_point_(zero_point),
+        clamp_min_(clamp_min),
+        clamp_max_(clamp_max),
+        scale_double_(scale),
+        zero_point_double_(zero_point),
+        clamp_min_double_(clamp_min),
+        clamp_max_double_(clamp_max),
+        storage_bit_width_(storage_bit_width),
+        is_signed_(is_signed),
+        round_mode_(APFloat::rmNearestTiesToAway) {}
+
+  UniformQuantizedValueConverter(double scale, double zero_point,
+                                 const APFloat& clamp_min,
+                                 const APFloat& clamp_max,
+                                 uint32_t storage_bit_width, bool is_signed)
+      : scale_(scale),
+        zero_point_(zero_point),
+        clamp_min_(clamp_min),
+        clamp_max_(clamp_max),
+        scale_double_(scale),
+        zero_point_double_(zero_point),
+        clamp_min_double_(clamp_min.convertToDouble()),
+        clamp_max_double_(clamp_max.convertToDouble()),
+        storage_bit_width_(storage_bit_width),
+        is_signed_(is_signed),
+        round_mode_(APFloat::rmNearestTiesToAway) {}
+
+  virtual APInt quantizeFloatToInt(APFloat expressed_value) const {
+    // This function is a performance critical code path in quantization
+    // since it runs for each single float parameter value.
+
+    // Specialize f32->u8/i8 case to optimize performance.
+    if (&expressed_value.getSemantics() == &APFloat::IEEEsingle() &&
+        storage_bit_width_ == 8 &&
+        round_mode_ == llvm::APFloatBase::rmNearestTiesToAway) {
+      return quantizeF32ToInt8(expressed_value);
+    }
+
+    bool lossy;
+    expressed_value.convert(scale_.getSemantics(), round_mode_, &lossy);
+    // fixed_point = clamp(clamp_min, clamp_max, (
+    //   roundHalfToEven(expressed / scale) + zero_point))
+    APFloat scaled = (expressed_value / scale_);
+    scaled.roundToIntegral(round_mode_);
+    scaled.add(zero_point_, round_mode_);
+    APFloat fixed_point = llvm::minimum(scaled, clamp_max_);
+    fixed_point = llvm::maximum(fixed_point, clamp_min_);
+
+    llvm::APSInt result(storage_bit_width_, !is_signed_);
+    fixed_point.convertToInteger(result, round_mode_, &lossy);
+
+    return std::move(result);
+  }
+
+  int64_t quantizeFloatToInt64(APFloat expressed_value) const {
+    const APInt q_value = quantizeFloatToInt(std::move(expressed_value));
+    return is_signed_ ? q_value.getSExtValue() : q_value.getZExtValue();
+  }
+
+  virtual ~UniformQuantizedValueConverter() = default;
+
+ private:
+  // An optimized implementation to quantize f32 to i8/u8 with C++ native
+  // arithmetic.
+  virtual APInt quantizeF32ToInt8(const APFloat& expressed_value) const {
+    assert(&expressed_value.getSemantics() == &APFloat::IEEEsingle());
+    assert(storage_bit_width_ == 8);
+    assert(round_mode_ == llvm::APFloatBase::rmNearestTiesToAway);
+
+    const float real_value = expressed_value.convertToFloat();
+
+    const double scaled = real_value / scale_double_ + zero_point_double_;
+    // Round to nearest integer with halfway cases rounded away from zero.
+    const double scaled_rounded = std::round(scaled);
+    const double clamped = std::min(std::max(scaled_rounded, clamp_min_double_),
+                                    clamp_max_double_);
+
+    uint64_t signless_result;
+    if (is_signed_) {
+      int64_t clamped_int = static_cast<int8_t>(clamped);
+      memcpy(&signless_result, &clamped_int, sizeof(clamped_int));
+    } else {
+      signless_result = static_cast<uint8_t>(clamped);
+    }
+    return APInt(storage_bit_width_, signless_result, /*isSigned=*/is_signed_);
+  }
+
+  // Keep both APFloat and double versions of the quantization parameters
+  // around since they will be used in generic and specialized arithmetic,
+  // respectively.
+  const APFloat scale_;
+  const APFloat zero_point_;
+  const APFloat clamp_min_;
+  const APFloat clamp_max_;
+
+  const double scale_double_;
+  const double zero_point_double_;
+  const double clamp_min_double_;
+  const double clamp_max_double_;
+
+  const uint32_t storage_bit_width_;
+  const bool is_signed_;
+  const llvm::APFloat::roundingMode round_mode_;
+};
+
+// An utility class to quantize an attribute by the per-axis quantization
+// parameters. The size of the quantization dim in the converted elements
+// attribute should match the size of of scales/zero_points vectors in the
+// quantization parameters.
+class UniformQuantizedPerAxisValueConverter {
+ public:
+  explicit UniformQuantizedPerAxisValueConverter(
+      quant::UniformQuantizedPerAxisType uniform_type)
+      : scales_(uniform_type.getScales()),
+        zero_points_(uniform_type.getZeroPoints()),
+        clamp_min_(static_cast<double>(uniform_type.getStorageTypeMin())),
+        clamp_max_(static_cast<double>(uniform_type.getStorageTypeMax())),
+        storage_bit_width_(uniform_type.getStorageTypeIntegralWidth()),
+        is_signed_(uniform_type.isSigned()),
+        quantization_dim_(uniform_type.getQuantizedDimension()) {
+    assert(isa<FloatType>(uniform_type.getExpressedType()));
+    assert(uniform_type.getStorageType().isSignlessInteger());
+    assert(scales_.size() == zero_points_.size());
+  }
+
+  // Quantize an Attribute by the quantization parameters. Return nullptr if
+  // the conversion fails or the input array isn't an ElementsAttr.
+  ElementsAttr convert(Attribute real_value);
+
+ private:
+  // Quantize an DenseFPElementsAttr by the quantization parameters.
+  DenseElementsAttr convert(DenseFPElementsAttr attr);
+
+  // Get a uniform converter for the index-th chunk along the quantizationDim.
+  // All the elements in this chunk is quantized by the returned converter.
+  UniformQuantizedValueConverter getPerChunkConverter(int index) const {
+    return UniformQuantizedValueConverter(scales_[index], zero_points_[index],
+                                          clamp_min_, clamp_max_,
+                                          storage_bit_width_, is_signed_);
+  }
+
+  const ArrayRef<double> scales_;
+  const ArrayRef<int64_t> zero_points_;
+  const APFloat clamp_min_;
+  const APFloat clamp_max_;
+  const uint32_t storage_bit_width_;
+  const bool is_signed_;
+  int32_t quantization_dim_;
+};
+
+}  // namespace mlir::quant::ir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_IR_UNIFORMSUPPORT_H_