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

  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
    format by inserting `TransposeOp`s to input, filter, and output tensors.

// LHS transpose permutation must be a NCHW->NHWC permutation.
template <typename OpT>
void DeferRhsTransposeForBinaryOp(OpT op, PatternRewriter& rewriter) {
  auto transpose_op = cast<TransposeOp>(op.getOperand(0).getDefiningOp());
  Value lhs_pre_transpose = transpose_op.getOperand();

  // NCHW -> NHWC for the right-hand side, to match the operand's shape.
  Value rhs = op.getOperand(1);

inline constexpr StringRef kQuantizationMethodAttr = "_quantization_method";

// Permutation from the NHWC tensor format to NCHW. This is an inverse
// permutation of `kNchwToNhwcPermutation`.
inline constexpr std::array<int64_t, 4> kNhwcToNchwPermutation = {0, 3, 1, 2};

// Permutation from the NCHW tensor format to NHWC. This is an inverse
// permutation of `kNchwToNhwcPermutation`.

    if (data_format == "NCHW") {
      x = Transpose(scope, x, {0, 2, 3, 1});
      grad_y = Transpose(scope, grad_y, {0, 2, 3, 1});
    } else if (data_format == "NCDHW") {
      x = Transpose(scope, x, {0, 2, 3, 4, 1});
      grad_y = Transpose(scope, grad_y, {0, 2, 3, 4, 1});
    }

    StringPiece target_data_format;
    if (data_format == "NCHW" || data_format == "NHWC") {

def TF_ConvnetDataFormatAttr : StringBasedAttr<
    CPred<"$_self.cast<StringAttr>().getValue() == \"NHWC\" || " #
          "$_self.cast<StringAttr>().getValue() == \"NCHW\"">,
    "'NHWC' or 'NCHW' convnet data format">;

//===----------------------------------------------------------------------===//
// Type attributes

// A derived attribute that returns the size of `idx`-th ODS-declared variadic

# CHECK: fake/user/code/file_C.py:27:1: error: 'tf.Conv2D' op attribute 'data_format' failed to satisfy constraint: 'NHWC' or 'NCHW' convnet data format

node {
  name: "input"
  op: "Placeholder"
  attr {
    key: "dtype"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "shape"
    value {
      shape {

// CHECK: stablehlo.reduce_window

// -----

// [No change] Tests that the transpose of the input of
// `stablehlo.reduce_window` is NOT deferred to the result, when the transpose
// isn't `[0, 3, 1, 2]` (i.e. NCHW->NHWC).

// CHECK-LABEL: reduce_window_max_activation_transpose_with_padding
func.func @reduce_window_max_activation_transpose_with_padding(%arg0: tensor<16x16x4x1xf32>) -> tensor<1x4x8x8xf32> {

  // The result of the conv must be the first input to BiasAdd to be fusable.
  %1 = "tf.BiasAdd"(%0, %arg0) {data_format = "NCHW"} : (tensor<*xf32>, tensor<128xf32>) -> tensor<*xf32>
  %2 = "tf.Elu"(%1) : (tensor<*xf32>) -> tensor<*xf32>
  %3 = "tf.Identity"(%2) : (tensor<*xf32>) -> tensor<*xf32>
  func.return %3 : tensor<*xf32>
}

# CHECK: fake/user/code/file_C.py:27:1: error: 'tf.Conv2D' op attribute 'data_format' failed to satisfy constraint: 'NHWC' or 'NCHW' convnet data format
# CHECK: fake/user/code/file_D.py:28:1: note: called from
# CHECK: fake/user/code/file_E.py:29:1: note: called from
# CHECK: fake/user/code/file_F.py:30:1: note: called from

node {
  name: "input"