diff --git a/Project.toml b/Project.toml index 849b560..93ed50c 100644 --- a/Project.toml +++ b/Project.toml @@ -12,14 +12,17 @@ Republic = "27243419-9dde-4721-b67c-fd63626fea7f" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" [weakdeps] +cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd" cuTile = "0dea8319-8c4a-4662-a73d-20234d115b9a" [extensions] +cuDNNExt = "cuDNN" cuTileExt = "cuTile" [compat] Adapt = "4.6.1" Republic = "2.2.0" StaticArrays = "1.9.18" -cuTile = "0.3.2" +cuDNN = "6.2.1" +cuTile = "0.3.2, 1.0.0" julia = "1.10" diff --git a/ext/cuDNNExt.jl b/ext/cuDNNExt.jl new file mode 100644 index 0000000..a54272d --- /dev/null +++ b/ext/cuDNNExt.jl @@ -0,0 +1,16 @@ +module cuDNNExt + +import BitPacking +import cuDNN + +function cuDNN.checked_array_pointer(t::cuDNN.Tensor, a::BitPacking.NarrowArray) + parent(a) isa cuDNN.DenseCuArray || throw(ArgumentError( + "binding for $(t.name) needs GPU-backed storage, got $(typeof(parent(a)))")) + cuDNN.graph_dtype(eltype(a)) == t.dtype || throw(ArgumentError( + "binding for $(t.name) has eltype $(eltype(a)), expected $(t.dtype)")) + length(a) == prod(t.dims) || throw(DimensionMismatch( + "binding for $(t.name) has $(length(a)) elements, expected $(prod(t.dims))")) + return pointer(parent(a)) +end + +end