Onednn: No format query interface for pooling/batch normalization

[CaoZhongZ]

The basic idea that pooling, batch normalization or relu doesn't support format query is that these interfaces suppose to be format agnostic, I get it. But currently we encountered pooling and batch normalization with 'nchw' go to ref path. It kind broke the promise. Any plan to solve this problem?

[CaoZhongZ]

Also we spotted that there is possible that in relu backward, grad_y had different format with x. Possibly that y was reorder out MKL-DNN's scope and they provided grad_y.

[vpirogov]

@4pao, we are extending the set of optimized implementations to support NCHW and NHWC formats across all primitives with exception of convolution and inner product. In particular, optimized batch normalization for NCHW, NHWC and NC formats will be promoted shortly. Pooling has optimized implementation for NCHW and work on NHWC version is in progress.

For the second scenario you describe, where src and dst have different layouts, MKL-DNN will indeed go to reference code and providing optimized implementations for this scenario does not seem feasible. Maybe in that case the library should just produce an error instead of running reference implementation.

[CaoZhongZ]

That will be great! Thank you!

[CaoZhongZ]

For the second scenario, I might add an example to explain it a little bit, like I have a small net:

convolution(x) -> relu(x)->dropout(x)
convolution_backward(gy) <- relu_backward(x, gy) <- dropout_backward(gy)

MKL-DNN does not have droptout so we use some code to do the job. In order to do that, we need to reorder y from nChw16c to nchw. And in backward path, dropout_backward provided us a 'gy' with format nchw. Then at relu_backward, x and gy had different format. In order to go fast, we have to reorder gy to nChw16c.

This solution is OK, yes, and we are using it. But consider large x and gy, Relu is memory bound in nature, so is reorder. We will go through two memory bound process, or we can fusion relu_backward to reorder, which might cut it into one. I'm just asking a little attention to this scenario, see if any easy way to gain about 1/3 performance relatively.

[vpirogov]

This makes sense actually. We do not see yet a manageable solution to support ops with different input and output layouts, as the number of combinations considering all the formats we support is very scary. If you have any ideas around the implementation approach I would be interested to discuss these.

[CaoZhongZ]

OK, I'll make my initial thought a little more detail and come back later. Basically, add fusion ops back to reorders, let the reorder be the driver.

[emfomenk]

Hi @4pao,

As you mentioned in second scenario user is to convert gy to the format of output of relu(x). This is a recommended way to handle the case. It solves 2 problems:

1. MKL-DNN has a fast implementation and
2. workspace (if any) would be in the right shape

But consider large x and gy, Relu is memory bound in nature, so is reorder. We will go through two memory bound process, or we can fusion relu_backward to reorder, which might cut it into one.

You will have to convert the data back to the blocked layout anyway. If not at relu_bwd step then at convolution_bwd step. So there should not be big performance difference when exactly that happens. Moreover, for things like pooling that shrinks the data of forward but expands it on backward earlier conversion is even desired.

Basically, add fusion ops back to reorders, let the reorder be the driver.

Would be interesting to here your ideas! One quick note -- we found it difficult to describe fusion for backward pass, since that typically changes the number of inputs/outputs which might very confusing for users... But anyway, please let us know what you think -- that would be great to discuss.

[vpirogov]

Batch norm optimizations are in master, see 3de46c50c97400ff47d546458b3e0425ed561751.

[CaoZhongZ]

Great! Thanks!

[vpirogov]

Pooling is optimized in cbd5f8ffe6fbb874012a35abe75c719193ae40b6. With that we have all the non-convolution primitives supporting NHWC and NCHW formats in addition to the blocked one. With that I'm closing this item.