Support torch.nn in slangpy

[wlewNV]

Issue: #535

[ccummingsNV]

Hi - I'd like to discuss what the design is here and what we're trying to achieve. I'm concerned about adding new python layers and torch specific wrappers, when we're trying to improve performance and avoid users having to think about torch - the goal of much of the torch integration work was to make it seemless.

I don't understand exactly what the requirements are to make slangpy 'torch.nn' compatible. It seems like this is a re-implementation of either 'Struct' or 'InstanceList' with an api to match that of torch.nn. Can we just make one of those 2 torch.nn compatible itself, or is there some requirement somewhere that torch.nn.Module is involved.

If the latter, how thin can we make that wrapper, baring in mind that we're currently working to scrape back microseconds off the cpu overhead of a dispatch.

[wlewNV]

Hey Chris,

I originally misinterpreted the issue. The comment "slangtorch seems to support using torch.nn directly" made me think we were trying to make a wrapper around the torch.nn.Module for user ease-of-use.

From your feedback, I tried using torch.nn directly in slangpy and ran into a bug where the torch.nn.Parameter wasn't recognized by slangpy's native code.

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.weight = nn.Parameter(torch.tensor(2.0, device="cuda"))
    
    def forward(self, x):
        return slangpy_module.my_func(self.weight, x)  # Failed!

The self.weight (nn.Parameter) was rejected with "Unsupported type" since nanobind's try_cast does exact type matching (for torch.Tensor), not isinstance checking.

I've reverted the original wrapper code and replaced it with a small C++ fix in unpack_arg() that handles Tensor subclasses.

@mkeshavaNV could you check to see if I'm aligned with your issue?

Thank you both!

[ccummingsNV]

Thanks @wlewNV

I think let's just hold off on this for a week or so, as our upcoming optimizations will probably intersect it a lot. We'll keep the PR here as all the tests will be valid, but I need to rewrite how the unpacking/signature checking works so it doesn't go via nanobind (this is apparently extremely expensive, especially on linux).

Once I start optimizing the tensor marshalling path, I'll merge your tests to my local branch and make sure they're accounted for as part of the work.