NCHW vs. NHWC

[mauricett]

Hey everyone!

I'm using JAX (with an Nvidia GPU), which wraps OpenXLA's convolution. If I understand the OpenXLA docs correctly, it expects the memory format for convolutions to be channel-first (NCHW), which is different from Nvidia's recommendation (NHWC). Could someone elaborate on this? Does this mean OpenXLA will perform sub-optimally?

Excuse the naive question... thanks.

[mauricett]

Okay, I've been testing more carefully and found that there's circumstances where OpenXLA can handle the NHWC format perfectly fine without introducing any extra transposes, but it seems to rely on the exact parameters of the convolution. With other parameters, I'm getting unnecessary extra transposes, but I don't know how much of this behaviour might be due to JAX.

Here's an example where a NHWC convolution does 4 transposes in total:

import jax.numpy as jnp
import jax


dtype = 'bfloat16'
out_dtype = 'float32'

conv_args = {'window_strides': (1, 1),
            'padding': 'same',
            'precision': dtype,
            'preferred_element_type': out_dtype,
            'dimension_numbers': ('NHWC', 'HWIO', 'NHWC')}

key = jax.random.key(0)
x = jax.random.normal(key, (2048, 8, 8, 128), dtype)
W = jax.random.normal(key, (3, 3, 128, 128), dtype)

f = jax.jit(lambda i: jax.lax.conv_general_dilated(i, W, **conv_args))
y = f(x).block_until_ready()

jax.profiler.start_trace("/tmp/tensorboard")
for i in range(10):
    y = f(x).block_until_ready()
jax.profiler.stop_trace()