// 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

   (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>,

    return %1 : tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
  }
}

// -----

// Merge fusion with dequantize for no activation case.

module attributes {tf_saved_model.semantics} {
  // CHECK-LABEL: func.func private @merge_no_act_fusion
  func.func private @merge_no_act_fusion(%arg0: tensor<1x4xf32>) -> tensor<1x3xf32> {

def TFR_TFRQuantActRangeOp : TFR_Op<"quant_act_range", [Pure]> {
  let description = [{
   The `quant_act_range` returns the a pair of integers to indicate the fixed
   range for the fused activation `act` with the quantization defined by the
   `scale` and `zero point`. Currently, the allowed activations are
   `NONE`, `RELU`, `RELU6` and `RELU_N1_TO_1`.

    Example:

    ```mlir