Implement reverse lookup (Ptr->Tuple) for CUDNN descriptors.

[RomeoV]

This fixes #1947.

[RomeoV]

Here's a small script that shows the loading and saving

import Pkg;
Pkg.activate("@CUDAPerfCachingTest")
# Setup:
# I don't think you can link to the cuDNN.jl module within CUDA.jl directly, so
# you'll have to clone github.com/romeov/CUDA.jl and then link
# Pkg.develop(path="<local>/romeov/CUDA.jl/lib/cudnn")
# Also
# Pkg.add("Flux")
# Pkg.add("JLD2")
#
# Execute e.g. with `julia caching_test.jl save` or `julia caching_test.jl load` or just `julia caching_test.jl`
using Flux, JLD2
import Flux.Zygote: gradient

function load_conv_caches!(; cudnn_mod::Module=Flux.cuDNN, filename="/tmp/conv_cache.jld2")
    @info "Loading conv_cache."
    conv_data_cache = JLD2.load(filename, "conv_data_cache");
    push!(cudnn_mod.cudnnConvolutionBwdDataAlgoPerfCache,
          conv_data_cache...)
    conv_filter_cache = JLD2.load(filename, "conv_filter_cache");
    push!(cudnn_mod.cudnnConvolutionBwdFilterAlgoPerfCache,
          conv_filter_cache...)
end

function save_conv_caches(; cudnn_mod::Module=Flux.cuDNN, filename="/tmp/conv_cache.jld2")
    @info "Storing conv_cache."
    JLD2.save(filename,
              "conv_data_cache", cudnn_mod.cudnnConvolutionBwdDataAlgoPerfCache,
              "conv_filter_cache", cudnn_mod.cudnnConvolutionBwdFilterAlgoPerfCache,
    )
end



if "load" in ARGS
    load_conv_caches!()
end

model = Chain(Conv((3, 3), 3=>64, relu; pad=SamePad()),
              Conv((3, 3), 64=>32, relu),
              GlobalMeanPool(),
              Flux.flatten,
              Dense(32=>1))

x = rand(Float32, 32, 32, 3, 7);

let x = gpu(x),
    model = gpu(model),
    ps = Flux.params(model)

  t0 = time()
  gradient(ps) do
      model(x) |> sum
  end
  @info "done in $(time() - t0) seconds :)"
end;
("load" in ARGS) && @show length(Flux.cuDNN.cudnnConvolutionBwdDataAlgoPerfCache)

if "save" in ARGS
    save_conv_caches()
end

[ToucheSir]

Instead of recreating the cache in reversed form and searching it every time (expensive!), CUDA.jl provides functions for pulling out the info from a descriptor (cheap!). See

CUDA.jl/lib/cudnn/src/tensor.jl

Lines 49 to 72 in 864ec5e

	function cudnnGetTensorDescriptor(d::cudnnTensorDescriptor)
	nbDimsRequested = CUDNN_DIM_MAX
	dataType = Ref{cudnnDataType_t}(CUDNN_DATA_FLOAT)
	nbDims = Ref{Cint}(0)
	dimA = Array{Cint}(undef, CUDNN_DIM_MAX)
	strideA = Array{Cint}(undef, CUDNN_DIM_MAX)
	cudnnGetTensorNdDescriptor(d, nbDimsRequested, dataType, nbDims, dimA, strideA)
	T = juliaDataType(dataType[])
	D = (dimA[nbDims[]:-1:1]...,)
	S = (strideA[nbDims[]:-1:1]...,)
	return T,D,S
	end
	function cudnnGetFilterDescriptor(d::cudnnFilterDescriptor)
	nbDimsRequested = CUDNN_DIM_MAX
	dataType = Ref{cudnnDataType_t}(CUDNN_DATA_FLOAT)
	format = Ref{cudnnTensorFormat_t}(CUDNN_TENSOR_NCHW)
	nbDims = Ref{Cint}(0)
	dimA = Array{Cint}(undef, CUDNN_DIM_MAX)
	cudnnGetFilterNdDescriptor(d, nbDimsRequested, dataType, format, nbDims, dimA)
	T = juliaDataType(dataType[])
	D = (dimA[nbDims[]:-1:1]...,)
	return T,D,format[]
	end

. You'll have to write equivalent functions for some of the conv-specific descriptors, but it should be quite straightforward.

[RomeoV]

Done.

[RomeoV]

The descriptors as they are still have some Cenum types in them, which we could convert to julia Ints or something if we run into serialization trouble.

[maleadt]

This is still marked WIP; anything to do here @RomeoV @ToucheSir?

[ToucheSir]

From my end no, didn't even notice the PR title still had WIP.

[ToucheSir]

@RomeoV whoops, I think I missed this line. It should be get(cudnnConvolutionBwdFilterAlgoPerfCache, key, nothing), right?

[RomeoV]

Good catch, thanks. Opened another PR with that one-line change.