[XLA:GPU/CuDNN] Add support for 1x1 window reversal
CuDNN supports both convolution (CUDNN_CONVOLUTION) and cross correlation
(CUDNN_CROSS_CORRELATION). However, only the latter was hooked up, causing
Tensorflow error: Status: Hit a case for convolution that is not implemented on GPU.
for convolutions of the first kind (corresponding to convolutions with both
window dimensions reversed at the HLO level). Reversing the dimensions (i.e.
doing convolutions rather than cross correlations) is the default behavior
for the Flux.jl ML framework, so it's easy to hit this error trying to run
pre-existing Flux models through the Julia:XLA->XLA:GPU compilation path.
Plumb through the reversal option to CuDNN to make this pattern work. The
same HLO already works fine against the CPU and TPU backends.
Loading
Please sign in to comment