Onednn: relu_backward with y as input instead of x

[eric-haibin-lin]

Let say y = relu(x), in the backward pass to calculate dx, currently mkl-dnn takes dy and x as the input for backward computation.

But for relu, to calculate the backward pass we can just pass y and dy to calculate dx. This way during forward pass relu can be done inplace (y can sharing the same memory as x), which saves significant amount of memory.

I wonder if we can update/add relu_backward API to accept y as the input instead of x. Mathematically it should be fine since for relu it only checks the non-zeros but it'll be great to get confirmation from mkl-dnn side.

[emfomenk]

Hi @eric-haibin-lin,

You are right, ReLU is kinda special from that perspective.
In fact MKL-DNN supports in-place ReLU for both forward and backward passes.
That works because as you said there is no difference whether you pass y or x to the backward pass.

  y <- relu(x) // y and x might be the same memory
  ...
  dx <- d_relu(dy, y) // dx and dy also might be the same memory

Batch-normalization also supports in-place computations (on both fwd and bwd passes).

[eric-haibin-lin]

Thanks for your fast response!
I guess it's fine if I just pass y instead of x to the backward API then? I was not that comfortable because it was not mentioned in mkl-dnn doc, so I don't know if this is a valid operation

[emfomenk]

Yes, passing y instead of x is absolutely fine.

Yeah... we should've mentioned that in the documentation, agree. So far the only note we have is in the example, bullet 11. But the example is about inference only...

[eric-haibin-lin]

Thanks. Could somebody add this (passing y instead of x to backward_relu and inplace optimzation) to a more obvious place in the mkl-dnn documentation? That will be very helpful to others :)

[emfomenk]

Sure, I will take care of that.
Thanks for bringing this!

[emfomenk]

I still have to update documentation.
Let the bug be open for a while...