Add support for tt.dot using 3d tensors

[Daniel (Dasor)]

Really simple change to emit linalg.batch_matmul instead of linalg.matmul if the input to tt.dot is a 3d tensor.

[Ziliang Zhang (ziliangzl)]

Hi Daniel (@Dasor), I'm just curious if there would be a Triton kernel that could actually utilize a 3d tt.dot?

[Daniel (Dasor)]

Hi Ziliang Zhang (@ziliangzl) yes there are, I added support for it because this test from triton-cpu was failing:

https://github.com/triton-lang/triton-cpu/blob/main/python/test/unit/language/test_core.py#L4203

[Ziliang Zhang (ziliangzl)]

Oh I see, I honestly didn’t know that a 3D tt.dot could be implemented this way.
Thanks a lot for pointing me to the test, that’s really helpful!🙏

[Daniel (Dasor)]

May I have this PR and my other one review? Micah Weston (@red1bluelost)