Add block sparse linear and locally connected layers

[clane9]

No description provided.

[clane9]

I updated the implementation:

• Use pytorch native blocksparse tensors following this blog rather than triton blocksparse, which seems not yet stable.
• Represent weights in BlockSparseLinear as block-sparse tensor rather than dense tensor. This will save a lot of gpu memory.
• Change BlockSparseLocallyConnected interface to more closely match nn.Conv2d. Except remove support for padding and stride. For now we should be able to restrict to padding="same", stride=1 (@alismil what do you think?).
• Rewrite function for constructing local connectivity matrix to directly construct a sparse rather than dense matrix. Use vectorized operations rather than for loops for construction. This should save a lot of memory and run faster.
• Add support for multiple input and output channels and depthwise convolution. The channels axis can either be first or last. For depthwise convolution, first should be more efficient (more block sparsity).

[alismil]

For now we should be able to restrict to padding="same", stride=1 (@alismil what do you think?).

If we want to use this for All TNNs, to make the dimensions work out we have to use different strides for each layer

[clane9]

If we want to use this for All TNNs, to make the dimensions work out we have to use different strides for each layer

Ok good point. Then maybe we should add back the support you had for strides and padding to _sparse_local_connectivity. Hopefully it doesn't complicate things too much.

[alismil]

This is a bit restrictive, can we adapt this to also include even kernel size by doing something like this?

if kernel_size % 2 == 0:
    kernel_half_width = kernel_size // 2
    kernel_indices = torch.cartesian_prod(
        torch.arange(-kernel_half_width, kernel_half_width),
        torch.arange(-kernel_half_width, kernel_half_width),
    )
else:
    kernel_half_width = (kernel_size - 1) // 2
    kernel_indices = torch.cartesian_prod(
        torch.arange(-kernel_half_width, kernel_half_width + 1),
        torch.arange(-kernel_half_width, kernel_half_width + 1),
    )

[clane9]

Yes I think you're right, something like this would probably be better. Although it feels like it should be possible to make the code shorter.

More generally, it would probably be best to have exactly the same interface and behavior as Conv2d. What I have now takes a few shortcuts.