Stiletto.jl

Traces plain Julia array code into fused cuDNN graphs.

using Stiletto, CUDA

M, N, K = 640, 320, 480
A = CUDA.randn(Float32, K, M)
B = CUDA.randn(Float32, K, N)

function matmul_epilogue(a::AbstractMatrix, b::AbstractMatrix)
    sum(tanh.(transpose(a) * b / √K), dims=1)
end

C = @jit matmul_epilogue(A, B)   # 1×320 CuArray, 1 allocation

function matmul_epilogue!(c::AbstractMatrix, a::AbstractMatrix, b::AbstractMatrix)
    c .= matmul_epilogue(a, b)
end

@jit matmul_epilogue!(C, A, B)   # 0 allocations

Stiletto will try to run any graph supported by cuDNN, but not all graphs will run due limitations of cuDNN fusion and engine selection.

Installation

using Pkg
Registry.add(url="https://registry.jool.space")
Pkg.add("Stiletto")