This will silently alter the stride structure of out if upper == true. It is better be upper ? out.triu_() : out.tril_().

That's not the same. The kernel does decomposition in the lower part of the matrix. If you do out.triu_() instead of out.tril_ -> transpose, then you get the upper part of the matrix which isn't really the correct output.

Do you have some stride assumptions in the kernel, or is it stride-agnostic? If it is stride-agnostic, then the kernel could be run on the transposed variant.

It assumes that input is row major(contiguous)

out can be provided externally as column-major. What would happen in this case?

I printed data ptr inside the mps function and outside in python:

import torch

out = torch.rand(3, 3, 3, device="mps").permute(2, 1, 0)

x = torch.rand(3, 3, 3, device="mps")
x = x.mT @ x

data_ptr = out.data_ptr()
print(f"0x{data_ptr:x}")  # lowercase hex
torch.linalg.cholesky(x, out=out)
print(f"0x{out.data_ptr():x}")

Yields:

0x10a4d68d0
0x10fb19150
0x10a4d68d0

First one being print from python, 2nd one being before launching the kernel from C++ and 3rd one being again from python. So yeah confirmed

As per https://github.com/pytorch/pytorch/pull/146799/files#r1952464144, this is expected. Sorry for the confusion. But we should have issues when out is contiguous and upper=True it seems.

No issues from what I check:

import torch

out = torch.rand(3, 3, 3, device="mps").permute(2, 1, 0)

x = torch.rand(3, 3, 3, device="mps")
x = x.mT @ x

data_ptr = out.data_ptr()
print(f"0x{data_ptr:x}")  # lowercase hex
print(out.stride())
print(out.is_contiguous())
res1 = torch.linalg.cholesky(x, out=out, upper=True)
res2 = torch.linalg.cholesky(x.cpu(), out=out.cpu(), upper=True)
print(f"0x{out.data_ptr():x}")
print(out.stride())
torch.testing.assert_close(res1.cpu(), res2)

0x113f70cc0
(1, 3, 9)
False
0x114f3a510
0x113f70cc0
(1, 3, 9)

@Isalia20 , could you remove permute so that out is contiguous? In the Meta function, as per your modification, out is re-used only if it is contiguous.

Ah I see the issue now:

0x10bc7b840
(9, 3, 1)
True
0x10bc7b840
0x10bc7b840
(9, 1, 3)

Why is MPS different?

MPS Kernel assumes row-major layout for the matrix where it does the decomposition

Can the kernel be made row-major/col-major agnostic so as to be able preserve the consistency across backends?

I'll take a look ~next week to see if I can make it work for col-major so we don't need to make it row major for MPS only, but why do we want to preserve consistency across backends? Lot of ops on MPS use row major layout and require contiguous call on it before passing it to some MPS kernel

In linalg LAPACK seems like the source of truth, and it is written in Fortran where col-major is the standard layout :(

I believe we can re-use the kernel without that much code change (i.e. no need to make it stride-agnostic for now). In the Meta function we request C-contiguous when upper=False and F-contiguous when upper=True for the MPS. Then we only need to remove the line upper ? out.transpose_(...) : out (and probably replace it with out.tril_() : out.triu_(). Or something along these lines. Should resolve the issue for now with out, before the kernel is adapted for better memory accesses when in column-major mode...

I've tried it but I'm afraid it doesn't work. I'll address this in the followup PR with the kernel change for column major mode rather than going into the rabbit hole now for a temporary fix

Let us also check that info is the same since its behavior is altered?

output_cpu and output_mps is a tuple of L and info tensors so assertEqual is comparing both of them. Do you mean to add a separate test where info might be >1?

Yes, when erroring on non-psd inputs :)

I'll do it a bit later today and also adapt the error message

Added better error message