// CHECK: return %[[TRANSPOSE_1]]

// -----

// Tests that a `add(convolution(%activation, %weight), %bias)` pattern with the
// activation tensor of NCHW format and non-constant bias is converted to NHWC
// convolution, but without the deferred transpose for `stablehlo.add`.
// Transpose ops are inserted to the activation and output of
// `stablehlo.convolution`. The weight constants is transposed.

  ) -> tensor<?x2x2x1xi32> {
  // CHECK-DAG: %[[conv:.*]] = mhlo.convolution
  // CHECK-DAG: %[[combined:.*]] = chlo.broadcast_add %[[zp_offset:.*]], %[[bias:.*]]
  // CHECK-DAG: %[[result:.*]] = chlo.broadcast_add %[[conv]], %[[combined]]
  // CHECK: return %[[result]]
  %0 = mhlo.convolution(%lhs, %rhs)
      dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f],

  // Iterate through block argument and its convolution users. Space to depth
  // transform will be applied only if all the below conditions are satisfied:
  //  1. All the users of the block argument will lead to convolutions;
  //  2. block_size of for the space to depth transform for these convolutions
  //     are the same;
  //  3. block_size of for the space to depth transform for these convolutions
  //     are larger than 1.

using ::stablehlo::quantization::QuantizedType;
using ::stablehlo::quantization::WeightOnlyPtq;

// Inserts quantization parameters of weights for weight-only quantization and
// dynamic range quantization of `stablehlo.convolution` and
// `stablehlo.dot_general`.
class InsertWeightParamPass
    : public impl::InsertWeightParamPassBase<InsertWeightParamPass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InsertWeightParamPass)

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

def NchwConvolutionToNhwcPass : Pass<"stablehlo-nchw-convolution-to-nhwc", "mlir::func::FuncOp"> {
  let summary = "Converts stablehlo.convolution op of NCHW format to -> NHWC.";
  let description = [{
    Matches `ConvolutionOp`s with NCHW format and converts it to NHWC

namespace mlir {
namespace TFL {

// A dilated convolution can be emulated with a regular convolution by chaining
// SpaceToBatch and BatchToSpace ops before and after it:
//
//     SpaceToBatchND -> Conv2D -> BatchToSpaceND
//
// This method was common before Conv2D fully supported dilated convolution in
// TensorFlow. This transformation detects this "emulation", and replaces it
// with a true dilated convolution, eliminating the SpaceToBatch and

// CHECK-DAG: %[[CONST_0:.*]] = stablehlo.constant dense<[{{.*}}]> : tensor<2xf32>
// CHECK-DAG: %[[CONST_1:.*]] = stablehlo.constant dense<[{{.*}}]> : tensor<2x3x3x2xf32>
// CHECK-DAG: %[[CONV:.*]] = stablehlo.convolution(%[[ARG]], %[[CONST_1]]) {{.*}} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x2x2xf32>
// CHECK-DAG: %[[BROADCAST:.*]] = stablehlo.broadcast_in_dim %[[CONST_0]], dims = [3] : (tensor<2xf32>) -> tensor<1x3x2x2xf32>

  // once available.
  //
  // If set to true, enable channel-wise quantization for:
  //   * Convolution ops: When the attached `Method` also specifies per-channel
  //                      quantization.
  //   * Non-convolution ops: All
  //
  // Default value: true
  bool enable_per_channel_quantized_weight = 2 [deprecated = true];

// CHECK: @uniform_dequantize_0
}

// -----

// Tests a variant where there is no stablehlo.convert op in between the
// filter constant and the convolution op.
//
// `filter (f32) -> convolution`
//
// instead of:
//
// `filter (i8) -> convert (i8 -> f32) -> convolution`

module {
// CHECK-LABEL: quantized_conv_op_with_no_filter_convert
// CHECK-SAME: %[[ARG:.*]]: tensor<1x3x3x4xf32>

        %0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x4x4xf32>
        %1 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32>