Help setting up convolution in NNLib.jl

New to Julia
,
1

Hey there, new Julia user here trying to get acclimated.

As a warm-up example, I want to build a GPU accelerated heat-equation explicit time-stepper. I am trying to implement the discrete laplacian using a convolution in NNLib. However, I don’t understand the documentation and can’t find any clear examples that I can understand.

As a concrete example, I want to compute in a 10x10 grid with ghost cells, the convolution with kernel:

0 1 0
1 -4 1
0 1 0

The first step in setting up the convolution is constructing a DenseConvDims object. There is a “convenience wrapper” with signature

function DenseConvDims(x_size::NTuple{M}, w_size::NTuple{M};
                   stride=1, padding=0, dilation=1, flipkernel::Bool=false) where M

I’m not sure, but I think x,w are the dimensions of the source and target spaces. For this case, I believe I want the default stride, dilation and flipkernel. but I probably need padding=1 for the ghost cells. With this padding, I expect the target array to be smaller than the source space. So I try:

julia> DenseConvDims((11,11),(10,10),padding=1)
ERROR: DimensionMismatch("Input channels must match! (11 vs. 10)")

julia> DenseConvDims((11,11),(10,10),padding=0)
ERROR: DimensionMismatch("Input channels must match! (11 vs. 10)")

Ok, that didn’t work I’ll try to ignore the ghost cells for now.

julia> DenseConvDims((11,11),(11,11))
DenseConvDims: (11,) * () -> (11,), stride: () pad: (), dil: (), flip: false

This didn’t give any errors, but I don’t think this is correct. The stride, padding, dilation parameters appear unset, and the symbol (11,) * () -> (11,) doesn’t seem correct. I expect something like (11,11)*(3,3) -> (11,11) after I set the kernel, and before setting the kernel (11,11)*()->(11,11).

I don’t even know how to set up the kernel after this step.

Could somebody please help me with this example?

Thanks!