auto op_before_dequantize = original_dequantize_op.getOperand(0);

    // Create a new dequantize op that is propagated.
    rewriter.setInsertionPointAfter(user_op);
    TF::PartitionedCallOp new_dequantize_op =
        cast<TF::PartitionedCallOp>(rewriter.clone(*original_dequantize_op));

    // Skip the original dequant op and connect the op before dequantize to the
    // user op.
    user_op->setOperand(user_idx, op_before_dequantize);

    // Finally, use the quantization parameter to create the quantize and
    // dequantize ops, and insert them between the tf.FakeQuantWithMinMaxVarsOp
    // and its users.
    auto quantize = rewriter.create<quantfork::QuantizeCastOp>(
        tf_op.getLoc(), qtype.getValue(), input);
    auto dequantize = rewriter.create<quantfork::DequantizeCastOp>(
        tf_op.getLoc(), res_type, quantize.getResult());

// and tfl.dequantize pairs before tf.FakeQuant* being foled.
LogicalResult ConvertFakeQuantOps(func::FuncOp func, MLIRContext* ctx,
                                  bool use_fake_quant_num_bits) {
  OpBuilder builder(func);
  if (failed(UnwrapTFCustomOps(func, builder))) {
    return failure();
  }

  // Insert the tfl.quantize/tfl.dequantize ops after the tf.FakeQuant* ops to

include "mlir/IR/OpBase.td"
include "mlir/IR/PatternBase.td"
include "mlir/Dialect/Func/IR/FuncOps.td"
include "tensorflow/compiler/mlir/lite/ir/tfl_ops.td"

// Both Quantize and Dequantize ops have side effects, so we have to define
// patterns to remove dead ones after the quantization rewrite.
def : Pat<(TFL_QuantizeOp:$op $in, $qt), (replaceWithValue $in), [(HasNoUseOf:$op)]>;

}

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()";
}

RewritePatternSet GetHardwareRewritePatterns(MLIRContext* context,
                                             const std::string& hardware);

// Convert quantized ops to float, this will essentially insert dequantize &
// quantize pair around the op.
void ConvertQuantizedOpToFloat(func::FuncOp func, OpBuilder* builder);

// This will optimize the quantized ops -> float graph.

// `tfl.quantize` or `tfl.dequantize` ops. ui8, i8 and i16 are supported.
bool IsSupportedByTfliteQuantizeOrDequantizeOps(IntegerType storage_type) {
  if (storage_type.getWidth() == 8 ||
      (storage_type.isSigned() && storage_type.getWidth() == 16)) {
    return true;
  }
  LLVM_DEBUG(llvm::dbgs()
             << "Uniform quantize / dequantize op only supports ui8, i8 or "

    left as is for weight-only which means the weight is dequantized at runtime.

    For example, if the kernel does not support dynamic range quantization the
    graph will be converted into the following IR:

    %q_w = "tfl.pseudo_qconst"() {
         qtype = tensor<64x3x3x3x!quant.uniform<i8<-127:127>:f32, 1.000000e+00>>
    %w = "tfl.dequantize"(%q_w) :

// CHECK: %[[IDENTITY:.*]] = "tf.Identity"(%cst_0) : (tensor<200x100x300xi8>) -> tensor<200x100x300xi8>
// CHECK: %[[DEQUANTIZED:.*]] = "tf.PartitionedCall"(%[[IDENTITY]]) <{config = "", config_proto = "", executor_type = "", f = @composite_dequantize_uniform}> : (tensor<200x100x300xi8>) -> tensor<200x100x300xf32>

Jake Harmon <******@****.***> 1694027275 -0700