[Autodiff] Support backward in nn.Modules containing Helion kernels

[hinriksnaer]

When HELION_EXPERIMENTAL_DIFFERENTIABLE=1 is set, Helion kernels
automatically participate in the autograd graph. No manual
torch.autograd.Function boilerplate needed:

import helion
import helion.language as hl

@helion.kernel
def square_plus(x: Tensor) -> Tensor:
    out = torch.empty_like(x)
    for tile in hl.tile(x.size()):
        v = x[tile]
        out[tile] = v * v + v
    return out

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(128, 256)
        self.fc2 = nn.Linear(256, 32)

    def forward(self, x):
        x = self.fc1(x)
        x = square_plus(x)
        return self.fc2(x)

model = MyModel().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
x = torch.randn(16, 128, device="cuda")
target = torch.randn(16, 32, device="cuda")

for step in range(100):
    optimizer.zero_grad()
    loss = F.mse_loss(model(x), target)
    loss.backward()
    optimizer.step()

HELION_EXPERIMENTAL_DIFFERENTIABLE=1 python train.py

Kernel.__call__ detects tensors with requires_grad=True and routes
through a thin torch.autograd.Function wrapper. The backward pass
calls
helion.experimental.backward(), which auto-generates a backward Helion
kernel from the forward kernel's FX graph. Both forward and backward are
real Helion kernels compiled through the normal pipeline.