// RUN: stablehlo-quant-opt %s -stablehlo-process-nchw-tensor \
// RUN:   -split-input-file -verify-diagnostics | FileCheck %s

// Tests that a `convolution(%activation, %weight)` with the activation tensor
// NCHW format is converted to NHWC convolution. Transpose ops are inserted to
// the activation and output to match the function signature. The weight
// constant is transposed.

// CHECK-LABEL: nchw_conv
// CHECK-SAME: %[[ARG:.+]]: tensor<1x8x4x4xf32>

// RUN: stablehlo-quant-opt %s -stablehlo-defer-activation-transpose \
// RUN:   -split-input-file -verify-diagnostics | FileCheck %s

// Tests that an `add(transpose(arg0), arg1)` pattern is converted to
// `transpose(add(arg0, transpose(arg1)))`. The transpose in the activation is
// deferred to the output of `stablehlo.add` and an extra transpose op is
// inserted to the RHS to match the shape of the operand.

// CHECK-LABEL: add_with_activation_transpose

  (TF_MaximumOp (TF_ConstOp:$cst $cst_val), $input),
  (TF_MaximumOp $input, $cst)>;

def SwapMinimumOperands : Pat<
  (TF_MinimumOp (TF_ConstOp:$cst $cst_val), $input),
  (TF_MinimumOp $input, $cst)>;

// Relu1 activation is represented as a couple of Max and Min ops, The following
// patterns recognize and keep them as TF ops so they can be converted to the
// TFLite Relu1 op.
def MatchRelu1Pattern1 : Pat<
  (TF_MinimumOp:$min_op

    // conversion, then dequantized during inference.
    // Activation: f32, Weight: qi8, Bias: f32
    WEIGHT_ONLY = 1;

    // Apply default dynamic range quantization. Quantized tensor value's
    // ranges are determined during graph runtime.
    // Activation: f32, Weight: qi8, Bias: f32
    POST_TRAINING_QUANTIZATION_DYNAMIC_RANGE = 2;

   (IsConstTensor $filter),
   (IsInt32ElementType $conv),
   (HasStaticShapeConstraint $filter),
   (HasStaticShapeAtDimsConstraint<"3"> $input)],
  [], (addBenefit 10)>;

// Convert Conv2D with hybrid inputs (f32 activation/int8 weight) to XlaConv
def ConvertTFConv2DToXLAConvOpWeightOnly : Pat<
  (TF_Conv2DOp:$conv
    $input,
    (TF_MulOp (TF_CastOp (TF_IdentityOp $filter), $truncate1), $scale),

//===----------------------------------------------------------------------===//
// Pattern rules for lifting ops with activation as functions
//===----------------------------------------------------------------------===//

def LiftConvWithRelu : Pat<
  (StableHLO_MaxOp:$res
    (StableHLO_ConvolutionOp $lhs, $rhs, $window_strides, $padding,

    there's any, and set it to True. The reason behind this decision is that
    generally activations of these ops show better accuracy with asymmetric
    input quantization so we want to deprecate symmetric activation quantization
    for those ops eventually.
    - Unlike to the old quantizer, per-channel quantization is supported for
    weight-only TransposeConvOp.
  }];

  let methods = [
    InterfaceMethod<

   (IsEinsumSupportedByXlaDotV2 $equation)],
  [], (addBenefit 5)>;

//===----------------------------------------------------------------------===//
// Pattern rules for lifting ops with bias and activation as functions
//===----------------------------------------------------------------------===//

multiclass LiftCompositeOpsWithActivation<Op ActivationOp, string ActivationName> {
  def LiftConvWith#ActivationOp : Pat<

  }];
  let dependentDialects = ["mlir::stablehlo::StablehloDialect"];
}

def DeferActivationTransposePass : Pass<"stablehlo-defer-activation-transpose", "mlir::func::FuncOp"> {
  let summary = "Merges stablehlo.transpose for activations.";
  let description = [{
    Defers activation transposes (e.g. LHS of `stablehlo.add`) to the output and
    optionally inserts `stablehlo.transpose`s to match the shape of operands.

    %mul = "tf.Mul"(%cast, %scale) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
    func.return %mul : tensor<*xf32>
  }

  // Requantizes and clips to the range of quantized type if there is no specific activation.
  func.func private @internal_requantize_no_activation_fn(%accumulation : tensor<*xi32>,
                         %input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,