Repository: tensorflow
Issue: Convolution much slower when activating JIT compile
Labels: stat:awaiting tensorflower,type:performance,comp:tf.function,TF 2.15

### Issue type

Performance

### Have you reproduced the bug with TensorFlow Nightly?

No

### Source

binary

### TensorFlow version

2.14.0 and 2.15.0

### Custom code

Yes

### OS platform and distribution

Ubuntu 23.04

### Mobile device

_No response_

### Python version

3.10

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

The code at https://github.com/ponder-lab/samples/blob/bc487180eec5b576054e04dc6add2bc2e26e7a3e/tensorflow_custom_convolution/tensorflow_custom_convolution.py runs much slower when using the `jit_compile=True` argument to `tf.function`. With either `jit_compile=False` or no argument at all, it runs much faster.

### Standalone code to reproduce the issue

```shell
#!/usr/bin/env python

import time
import numpy as np
import tensorflow as tf
print("TensorFlow version:", tf.__version__)
assert(tf.__version__ == "2.15.0")

class MyConv2D(tf.keras.layers.Layer):

    def __init__(self, filters, kernel_size, strides=(1, 1), padding="VALID",
                 *args, **kwargs):
        super().__init__(*args, **kwargs)
        
        self.filters = filters
        self.kernel_size = kernel_size
        self.strides = strides
        self.padding = padding

    def build(self, input_shape):

        weights_shape = (self.filters,
                         self.kernel_size[0] * self.kernel_size[1] * input_shape[-1])
        w = tf.random.normal(weights_shape, stddev=0.01)
        self.w = self.add_weight(shape=w.shape, trainable=True)
        self.w.assign(w)

        b = tf.zeros((self.filters))
        self.b = self.add_weight(shape=b.shape, trainable=True)
        self.b.assign(b)

        super().build(input_shape)

    @tf.function(jit_compile=True)
    def call(self, inputs):

        # Extract patch
        #  patches dim: (batch, h, w, kernel_size[0]*kernel_size[1]*ich)
        patch_sizes = (1,) + self.kernel_size + (1,)
        patch_strides = (1,) + self.strides + (1,)
        patches = tf.image.extract_patches(inputs, sizes=patch_sizes,
                                           strides=patch_strides,
                                           rates=(1, 1, 1, 1),
                                           padding=self.padding)

        # Multiply weights
        mul =  self.w * tf.expand_dims(patches, 3)
        madd = tf.reduce_sum(mul, axis=-1)
        maddb = madd + self.b

        return maddb

def get_model_normal(input_shape):
    input_node = tf.keras.Input(input_shape)
    x = input_node

    x = tf.keras.layers.Conv2D(32, (3, 3), padding="same")(x)
    x = tf.keras.layers.Activation("relu")(x)
    x = tf.keras.layers.MaxPooling2D()(x)

    x = tf.keras.layers.Conv2D(32, (3, 3), padding="valid")(x)
    x = tf.keras.layers.Activation("relu")(x)
    x = tf.keras.layers.MaxPooling2D()(x)

    x = tf.keras.layers.Conv2D(32, (3, 3), padding="same")(x)
    x = tf.keras.layers.Activation("relu")(x)
    x = tf.keras.layers.MaxPooling2D()(x)

    x = tf.keras.layers.Conv2D(32, (3, 3), padding="valid")(x)
    x = tf.keras.layers.Activation("relu")(x)
    x = tf.keras.layers.Flatten()(x)

    x = tf.keras.layers.Dense(10)(x)
    x = tf.keras.layers.Activation("softmax")(x)

    output_node = x
    return tf.keras.models.Model(input_node, output_node)

def get_model_myconv(input_shape):
    input_node = tf.keras.Input(input_shape)
    x = input_node

    x = MyConv2D(32, (3, 3), padding="SAME")(x)
    x = tf.keras.layers.Activation("relu")(x)
    x = tf.keras.layers.MaxPooling2D()(x)

    x = MyConv2D(32, (3, 3), padding="VALID")(x)
    x = tf.keras.layers.Activation("relu")(x)
    x = tf.keras.layers.MaxPooling2D()(x)

    x = MyConv2D(32, (3, 3), padding="SAME")(x)
    x = tf.keras.layers.Activation("relu")(x)
    x = tf.keras.layers.MaxPooling2D()(x)

    x = MyConv2D(32, (3, 3), padding="VALID")(x)
    x = tf.keras.layers.Activation("relu")(x)
    x = tf.keras.layers.Flatten()(x)

    x = tf.keras.layers.Dense(10)(x)
    x = tf.keras.layers.Activation("softmax")(x)

    output_node = x
    return tf.keras.models.Model(input_node, output_node)

(x_train, y_train), _ = tf.keras.datasets.mnist.load_data()

x_train = x_train.astype(np.float32)[..., None] / 255

model = get_model_myconv(x_train[0].shape)
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
model.summary()

model(x_train[:2])

myconv_start = time.time()

model.fit(x_train, y_train,
          batch_size=100, epochs=1, verbose=1,
          validation_split=0.1)

myconv_end = time.time()

del model
model = get_model_normal(x_train[0].shape)
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
model.summary()

model(x_train[:2])

normal_start = time.time()

model.fit(x_train, y_train,
          batch_size=100, epochs=1, verbose=1,
          validation_split=0.1)

normal_end = time.time()

print("Normal convolution: %f[sec]" % (normal_end - normal_start))
print("My convolution: %f[sec]" % (myconv_end - myconv_start))
```


### Relevant log output

_No response_