Nikita Shulga
11bb94b7ea
[MPSInductor] Fix index generation for transpose ( #143973 )
...
Alas, PythonPrinter would not work here, not would CppPrinter, so start building MetalPrinter.
`pytest test/inductor/test_torchinductor.py -k _mps` score is 474 failed, 277 passed, 32 skipped
Before this change:
`pytest test/inductor/test_torchinductor.py -k _mps` reported 506 failed, 245 passed, 32 skipped
Pull Request resolved: https://github.com/pytorch/pytorch/pull/143973
Approved by: https://github.com/jansel
ghstack dependencies: #143948 , #143949
2024-12-31 02:04:50 +00:00
Nikita Shulga
d8c3900d80
[Inductor] Implement primitive Metal compiler ( #143893 )
...
Still work in progress, only works for element wise operations. Current implementation could be used to turn something like
```python
def f(x):
return x[:,::2].sin() + x[:, 1::2].cos()
```
into the following shader
```python
# Topologically Sorted Source Nodes: [sin, cos, add], Original ATen: [aten.sin, aten.cos, aten.add]
# Source node to ATen node mapping:
# add => add
# cos => cos
# sin => sin
# Graph fragment:
# %sin : [num_users=1] = call_function[target=torch.ops.aten.sin.default](args = (%slice_2,), kwargs = {})
# %cos : [num_users=1] = call_function[target=torch.ops.aten.cos.default](args = (%slice_4,), kwargs = {})
# %add : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%sin, %cos), kwargs = {})
mps_lib = torch.mps._compile_shader("""
kernel void kernel_0(
device float* out_ptr0,
constant float* in_ptr0,
uint xindex [[thread_position_in_grid]]
) {
int x0 = xindex;
auto tmp0 = in_ptr0[2*x0];
auto tmp1 = metal::precise::sin(tmp0);
auto tmp2 = in_ptr0[2*x0 + 1];
auto tmp3 = metal::precise::cos(tmp2);
auto tmp4 = tmp1 + tmp3;
out_ptr0[x0] = static_cast<float>(tmp4);
}
""")
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/143893
Approved by: https://github.com/jansel
ghstack dependencies: #143891 , #143892
2024-12-28 06:58:32 +00:00
