Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implemented Cholesky on CPU #1026 #1119

Merged
merged 4 commits into from
May 17, 2024
Merged

Conversation

arn4
Copy link
Contributor

@arn4 arn4 commented May 14, 2024

Proposed changes

Implemented Cholesky Decomposition for real matrix on CPU, as requested in #1026.

There is nothing done for the GPU.

Checklist

Put an x in the boxes that apply.

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

Copy link
Member

@angeloskath angeloskath left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks fantastic! Left some tiny nitpicky comments for small typos and such.

mlx/backend/common/cholesky.cpp Outdated Show resolved Hide resolved
mlx/backend/common/cholesky.cpp Outdated Show resolved Hide resolved
mlx/backend/common/cholesky.cpp Outdated Show resolved Hide resolved
has more than two dimensions, the Cholesky decomposition is computed for each matrix
in the last two dimensions of ``a``.

If the input matrix is not symmetric positive semi-definite, behaviour is undefined.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It currently crashes. Maybe undefined is indeed better. We could just not throw in that case.

Copy link
Contributor Author

@arn4 arn4 May 16, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, it does crash if the matrix is not positive semidefinite. The behavior I wanted was

  • raise when the matrix is not positive semidefinite, since it is directly detected by Lapack
  • undefined if the matrix is not symmetric, since adding a check will come an unnecessary overhead.

Maybe I should specify it in the doc, or do you still prefer not raising at all?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I changed it to do nothing and added a comment to add some error checking if we can throw from the implementation in the future. Otherwise, potentially crashing the program in an unrecoverable way based on the values of the input matrix seems an extremely bad idea to me.

tests/linalg_tests.cpp Outdated Show resolved Hide resolved
mlx/backend/metal/primitives.cpp Outdated Show resolved Hide resolved
if (upper) {
std::fill(matrix, matrix + row, 0);
} else {
std::fill(matrix + row + 1, matrix + N, 0);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I missed it in the first review but here you were zeroing column-wise which is really inefficient especially for large matrices. I also use std::fill which may be faster if the compiler couldn't figure out that the previous loop was vectorizable.

@angeloskath angeloskath merged commit b3ec792 into ml-explore:main May 17, 2024
5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants