[REQUEST] Add Pickling support for Warp kernels

[eric-heiden]

Description

It is currently not possible to pickle Warp kernels.

Example:

import warp as wp
import pickle

@wp.kernel
def my_kernel(xs: wp.array(dtype=float), ys: wp.array(dtype=float)):
    tid = wp.tid()
    ys[tid] = xs[tid] * 2.0

pickled_kernel = pickle.dumps(my_kernel)

yields the error:

_pickle.PicklingError: Can't pickle <function my_kernel at 0x0000020199CA9000>: it's not the same object as __main__.my_kernel

Context

Being able to pickle kernels is necessary to save PyTorch models that contain Warp kernels. torch.save() uses pickle (or a user-provided pickling module) to store the models, which fails.