// it to NCHW before padding, and does all computations in NCHW (this is the
  // default setup for ResNet model trained in fp32 on GPU).
  //
  // To be able to use Tensor Cores on latest NVIDIA GPUs this model has to be
  // converted to NHWC data format.

  // Padding in spatial dimension (NCHW)

  //
  // Example:
  //   %1 = SomeOp(...)   {data_format = NHWC}
  //   %2 = Transpose(%1) {permutation = NHWC->NCHW}
  //   %3 = Op(%2)        {data_format = NCHW}
  //
  // To bypass %2 we have to change data format to shuffle data format from NCHW
  // to NHWC, which is the reverse of operand permutation (function argument).
  auto reverse_permutation =
      GetDataFormatPermutation(op->getDataFormat(), target_data_format);

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

  // MaxPool operand transpose must be folded into the op and MaxPool
  // must use NCHW data format with updated kernel size and strides.

  // CHECK: %[[RES_PERM:.*]] = "tf.Const"() <{value = dense<[0, 2, 3, 1]> : tensor<4xi32>}>
  // CHECK: %[[MAX_POOL:[0-9]*]] = "tf.MaxPool"(%arg0) <{data_format = "NCHW", ksize = [1, 1, 3, 3], padding = "SAME", strides = [1, 1, 2, 2]}> : (tensor<1x64x112x112xf32>) -> tensor<1x64x56x56xf32>

  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.

// RUN: tf-opt %s -tf-layout-assignment=force-data-format=NCHW -verify-diagnostics | FileCheck %s --dump-input=always

// IMPORTANT: In the following Conv2D tests tensor shapes do not match
// convolution parameters (stride, dilations, etc...). This test only verifies
// that changing convolution data layout will update all the attributes.

// CHECK-LABEL: func @transposeConv2D

  // NCHW->NHWC convolution conversion).
  pm.addNestedPass<func::FuncOp>(createFoldConstantTransposePass());
}

void RegisterPassPipelines() {
  static PassPipelineRegistration<> nchw_tensor_format_processing_pipeline(
      /*arg=*/"stablehlo-process-nchw-tensor",
      /*description=*/"Optimizes tensors with NCHW format.",
      AddProcessNchwTensorPasses);
}

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

  // CHECK: %[[RES_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%[[CONV2D]], %[[RES_PERM]])
  // CHECK: return %[[RES_TRANSPOSE]]

  %0 = "tf.Conv2D"(%input, %filter)
       {
         data_format = "NCHW",
         dilations = [1, 2, 3, 4],
         explicit_paddings = [1, 2, 3, 4, 5, 6, 7, 8],
         padding = "EXPLICIT",
         strides = [5, 6, 7, 8]

    reverse[permutation[i]] = i;
  }
  return reverse;
}

SmallVector<int64_t, 4> GetDataFormatPermutation(StringRef from, StringRef to) {
  if (from == "NHWC" && to == "NCHW") {
    return {0, 3, 1, 2};
  } else if (from == "NCHW" && to == "NHWC") {
    return {0, 2, 3, 1};
  } else {
    return {};
  }
}

// Shuffle elements in the `attr` according to the permutation. Optional