Add tests for PyTorch operators

[ScottTodd]

Could look at the Python test cases in https://github.com/pytorch/pytorch/tree/main/test, or lean more on tests in torch-mlir: https://github.com/llvm/torch-mlir/tree/main/test. There are also some tests in https://github.com/nod-ai/SHARK-TestSuite/tree/main/e2eshark/pytorch/operators.

Generally looking for proof that the work tracked in nod-ai/SHARK-ModelDev#347 has fully landed in IREE and is supported across the full range of targets.

[ScottTodd]

From @Groverkss here on Discord:

I've started using iree-turbine to generate tests. Writing IR for something like gather + attention fusion is somewhat hard for me, it's much easier to just write it in iree-turbine
This is the template i used:

import torch
import torch.nn as nn

import iree.turbine.aot as aot

class GatherAttention(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, q: torch.Tensor, k_mem: torch.Tensor, k_indices:
                torch.Tensor, v: torch.Tensor):
        k = torch.gather(k_mem, dim=2, index=k_indices)
        out = nn.functional.scaled_dot_product_attention(q, k, v)
        return out

model = GatherAttention()

q = torch.randn(2, 10, 4096, 64, dtype=torch.float16)
k_mem = torch.randn(2, 10, 4096, 64, dtype=torch.float16)
v = torch.randn(2, 10, 4096, 64, dtype=torch.float16)
k_indices = torch.arange(2 * 10 * 4096 * 64).reshape((2, 10, 4096, 64))

exported = aot.export(model, q, k_mem, k_indices, v)
exported.print_readable()

Would recommend generating tests like this because you can also check correctness with torch this way