Add BSR subclass +torch.compile and clean up superblock

[jcaip]

This PR adds in torch.compile support for block sparsity.

In a custom op, we create the sprase_bsr_tensor from the explicit crow_indices, col_indices, values tensors that are passed in to the custom op.

I also created a tensor subclass which holds these same values.

At dispatch, when we see a torch.nn.functional.linear call, we dispatch into our custom op torch.ops.blocksparse.linear, using the tensors stored in the subclass.

This will allow us to add a public API similar to semi_sparse_weight(), which I plan to do in a future PR.

This PR also cleans up the superblock prototype implementation, as there was a lot of repeated code, and also adds in kernel tuning for BSR.

For bfloat16 I see the following numbers, for a 1.23x gain:

New compile baseline: 63.431 ms
New compile + bsr: 53.514 ms 
New compile + bsr + tuning: 51.485 ms 

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/680

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit d14e320 with merge base accbdba ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[jerryzh168]

why not use

ao/torchao/utils.py

Line 51 in 5998389

def benchmark_model(model, num_runs, args=(), kwargs=None, device_type=None):

? we can add a warmup as well

[jcaip]

It was giving me different timings from my function. I'll try and debug the delta later, if I can find out whats causes it I'll switch over.