}
// Confirm that the `stablehlo.convolution` is not converted to `tfl.conv_2d`.
// CHECK-LABEL: convolution_upstream_srq_non_const_filter
// CHECK-SAME: %[[ARG:.+]]: tensor<1x3x3x4x!quant.uniform<i8:f32, 1.000000e+00:-100>>
// CHECK: stablehlo.convolution
// CHECK-NOT: tfl.conv_2d

// -----

// Tests that if the window padding contains values of 0, tfl.pad op is not

func.func @testConv2D(tensor<256x32x32x3xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>) -> tensor<256x32x32x16xf32> {
^bb0(%arg0: tensor<256x32x32x3xf32>, %arg1: tensor<16x3x3x3xf32>, %arg2: tensor<16xf32>):
  // CHECK: "tfl.conv_2d"(%arg0, %arg1, %arg2)

            FindUserOfType<TFL::QuantizeOp>(op) != nullptr);
  }
};

// Rewrites `stablehlo.convolution` into fused `tfl.conv_2d`.
// If available, fuse bias and activation adjacent to `stablehlo.convolution`.
// This RewritePattern rewrites both the following into `tfl.conv_2d` op:
//
// StableHLO Quantizer output:
//   * input: per-tensor qi8
//   * filter: per-channel qi8 (`quantization_dimension` = 3)

// with a 0-d constant, e.g. before this optimization,
//   %cst = arith.constant dense<1.0> : tensor<16x16x4xf32>
//   %0 = "tfl.conv_2d"...
//   %1 = "tfl.add"(%0, %cst) : (tensor<16x16x4xf32>, tensor<16x16x4xf32>)
// After this optimization:
//   %cst = arith.constant dense<1.0> : tensor<f32>
//   %0 = "tfl.conv_2d"...
//   %1 = "tfl.add"(%0, %cst) : (tensor<16x16x4xf32>, tensor<f32>)

      auto conv_result = rewriter.create<mhlo::ConvolutionOp>(
          conv_op.getLoc(), new_output_type, sliced_input, sliced_kernel,
          conv_op.getWindowStridesAttr(), conv_op.getPaddingAttr(),
          conv_op.getLhsDilationAttr(), conv_op.getRhsDilationAttr(),
          conv_op.getWindowReversalAttr(), conv_op.getDimensionNumbers(), 1, 1,
          conv_op.getPrecisionConfigAttr());

      conv_results.push_back(conv_result);

    in the flatbuffer.
  }];

  let arguments = (ins I32Attr:$buffer_index);

  let results = (outs AnyTensor:$output);
}

def TFL_Conv2DOp : TFL_ConvOp<"conv_2d", "Convolution", 0,
      [DeclareOpInterfaceMethods<InferTypeOpInterface>,
       DeclareOpInterfaceMethods<TFL_ArithmeticCount>,
       DynamicRangeQuantizedOpInterface]> {
  let arguments = (

        scale, offset = [1.0] * 2, [0.5] * 2
        mean, variance = scale, offset
        out = nn_ops.conv2d(
            q_input,
            q_filter,
            strides=[1, 1, 2, 1],
            dilations=[1, 1, 1, 1],
            padding='SAME',
            data_format='NHWC',
            name='sample/conv2d',
        )
        if has_bias:
          out = nn_ops.bias_add(out, bias, data_format='NHWC')

  %0 = "tf.Conv3D"(%arg0, %arg1) {padding = "SAME", strides = [1, 1, 1, 1, 1]} : (tensor<?x?x?x?x?xf32>, tensor<?x?x?x?x?xf32>) -> tensor<?x?x?x?x?xf32>
  func.return %0: tensor<?x?x?x?x?xf32>

  // CHECK-LABEL: conv3d_valid
  // CHECK:  %[[CST:.*]] = "tfl.no_value"() <{value}> : () -> none

        %device_launch = "tf_device.cluster_func"(%input,...) {func = @_func,...)
        return ...
      }
      func @_func(%input: tensor<2x224x224x3xf32>, %filter: tensor<7x7x3x64xf32>) {
        %6 = "tf.Conv2D"(%input, %filter)  {strides = [1, 2, 2, 1]}: (tensor<2x230x230x3xf32>, tensor<7x7x3x64xf32>) -> tensor<2x112x112x64xf32>
      }
    }
    ```

    The program will be transformed into:

    ```mlir

    return emitOptionalError(location, "invalid padding format provided");
  }

  // Output always have rank 4. All dimensions are initialized to
  // dynamic size and can be partially inferred.
  // TFL's conv2d is always NHWC format & the filter is OHWI.
  SmallVector<int64_t, 4> return_shape(4, ShapedType::kDynamic);
  return_shape[0] = input_ty.getDimSize(0);
  return_shape[3] = filter_ty.getDimSize(0);