%6 = "tf.Cast"(%5) {Truncate = false} : (tensor<1x100x100x1xi8>) -> tensor<1x100x100x1xf32>

  %7 = "tf.Sub"(%6, %cst_4) : (tensor<1x100x100x1xf32>, tensor<f32>) -> tensor<1x100x100x1xf32>
  %8 = "tf.Conv2D"(%7, %cst_3) {dilations = [1, 1, 1, 1], padding = "VALID", strides = [1, 1, 1, 1]} : (tensor<1x100x100x1xf32>, tensor<3x3x1x1xf32>) -> tensor<1x98x98x1xf32>

      return false;
    }
    return true;
  }
};
std::unique_ptr<TargetHardwareOperation> CreateConcatOp() {
  return std::make_unique<GpuConcatOp>();
}

// Currently used for these ops:
// tfl.conv_2d / tfl.depthwise_conv_2d / tfl.fully_connected
class GpuConvOp : public TargetHardwareOperation {
  double GetOpCost(mlir::Operation* op) const override {
    int64_t arithmetic_count;

      if (!input_type || !input_type.hasStaticShape()) {
        return false;
      }
      total_count += input_type.getNumElements();
    }
    *count = total_count;
    return true;
  }

  // For conv2d/depthwise_conv/fully_connected ops.
  // This algorithm actually comes from TOCO tooling_util.cc
  static bool GetArithmeticCountForConvAndFullyconnectedOp(mlir::Operation* op,

                         %input : tensor<*xf32>, %weight : tensor<*x!tf_type.qint8>,
                         %weight_scale : tensor<*xf32>, %weight_zp : tensor<*xi32>) -> tensor<*xf32>
      attributes {tf_quant.quantized_ops = ["Conv2D"]} {

    %out = "tf.UniformQuantizedConvolutionHybrid"(%input, %weight,
                           %weight_scale, %weight_zp) {
        Tlhs = "tfdtype$DT_FLOAT",
        Trhs = "tfdtype$DT_QINT8",

    if ((quantization_method_ == tensorflow::quantization::QuantizationMethod::
                                     METHOD_DYNAMIC_RANGE_INT8) &&
        (function_name.contains("batch_matmul") ||
         function_name.contains("conv3d"))) {
      call_op->removeAttr(kQuantTraitAttrName);
    }

    // TODO(b/270906404): Support weight-only gather for uniform quantized opset
    // in PTQ mode
    if (target_opset_ == OpSet::UNIFORM_QUANTIZED &&

  %1 = "tf.FakeQuantWithMinMaxArgs"(%arg0) {device = "", max = 2.000000e-01 : f32, min = -1.000000e-01 : f32, narrow_range = false, num_bits = 8 : i64} : (tensor<1x3x4x3xf32>) -> tensor<*xf32>
  %2 = "tf.Conv2D"(%1, %0) {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true} : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>

    %1 = "quantfork.stats"(%arg0) {layerStats = dense<[1.27501142, 149.824783]> : tensor<2xf32>} : (tensor<1x3x4x3xf32>) -> tensor<1x3x4x3xf32>

def DefinedByConv2D : Constraint<CPred<"llvm::isa_and_nonnull<mlir::TF::Conv2DOp>($0.getDefiningOp())">>;
// Checks if the value has only one user.
def HasOneUse : Constraint<CPred<"$0.hasOneUse()">>;

// If we see a Conv2D op followed by Mul, then multiply the filter
// with the value in Mul.
def FuseMulAndConv2D :
  Pat<(TF_MulOp:$mul (TF_Conv2DOp:$conv $input,
                          (Arith_ConstantOp:$filter F32ElementsAttr:$filter_value),

  %0 = "tf.Conv2D"(%output, %cst) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 2, 2, 1], use_cudnn_on_gpu = true}> {attr_map = "0:strides,1:use_cudnn_on_gpu,2:padding,3:explicit_paddings,4:dilations", device...

    %3 = "tf.Cast"(%2) {Truncate = false} : (tensor<960x960x3x512xi8>) -> tensor<960x960x3x512xi32>
    %4 = "tf.Sub"(%3, %arg5) : (tensor<960x960x3x512xi32>, tensor<512xi32>) -> tensor<960x960x3x512xi32>
    %5 = "tf.Conv2D"(%1, %4) {dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true} : (tensor<1x2240x2240x3xi32>, tensor<960x960x3x512xi32>) -> tensor<1x2240x1120x512xi32>