speedup and reduce memory usage in the first convolution.

The first convolution in many image models, including ResNet or ResNet-like, is
a (kernel=7, stride=2) 2D convolution. The input of the convolution is images,
which usually has RGB channels. The input of this first convolution is of shape
[batch\_size, height, width, 3] and the kernel size is [kernel\_size,
kernel\_size, 3, out\_channel]. Space to depth is to transform this first

  %2 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[o, i, 0, 1]->[b, f, 0, 1], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x4x4x8xf32>, tensor<8x8x3x3xf32>) -> tensor<1x8x4x4xf32>
  return %2 : tensor<1x8x4x4xf32>
}

// CHECK-NOT: stablehlo.transpose
// CHECK: %[[CONV:.+]] = stablehlo.convolution

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

// RUN: odml-to-stablehlo-opt %s -fuse-mhlo-convolution-pass -cse | FileCheck %s

// CHECK-LABEL: @fuseMulAndConv2D
// CHECK-SAME: %[[INPUT:[^:[:space:]]+]]
func.func @fuseMulAndConv2D(%input: tensor<1x256x256x3xf32>) -> (tensor<1x256x256x2xf32>) {
  // CHECK-DAG: %[[FILTER:.+]] = mhlo.constant dense<{{\[\[\[\[}}1.000000e+00, 2.000000e+00], [3.000000e+00, 4.000000e+00], [5.000000e+00, 6.000000e+00]]]]> : tensor<1x1x3x2xf32>

  %3 = stablehlo.constant dense<6.000000e+00> : tensor<f32>
  %4 = 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<?x28x28x1xf32>, tensor<3x3x1x16xf32>) -> tensor<?x28x28x16xf32>

module {
  func.func @main(%arg0: tensor<8x8x1x207xf32>, %arg1: tensor<3x3x16x207xf32>) -> tensor<16x8x8x1xf32> {

    rewriter.replaceAllUsesExcept(rhs, dq.getOutput(), dq);
  }
};

// Splits hybrid quantized `stablehlo.convolution` into `tfl.dequantize` and
// float `stablehlo.convolution` op. Weight tensor is transposed to match the
// filter tensor format for TFLite convolution.
// Legalization of float `stablehlo.convolution` op relies on existing passes
// for conversion of StableHLO -> MHLO -> TF -> TFL.

// CHECK-SAME: (tensor<1x2xf32>, tensor<2x3x!quant.uniform<i8:f32, 6.000000e-03>>) -> tensor<1x3xf32>
// CHECK: return %[[DOT]]

// -----

// Test that hybrid quantized convolution is produced when q/dq pair only exists
// for weight.

module attributes {tf_saved_model.semantics} {