Skip to content

add basic quaternion support from CUTLASS #987

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

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

simonbyrne
Copy link
Contributor

@simonbyrne simonbyrne commented Jun 5, 2025

This adds the most minimal required support for quaternions by aliasing quaternion from CUTLASS. As such it only works if CUTLASS support is enabled.

It mostly works out of the box: I tested basic arithmetic operators and math functions (e.g. exp).

Some open questions:

@cliffburdick
Copy link
Collaborator

Looks good! Was it too much work to break the dependency of CUTLASS?

@simonbyrne
Copy link
Contributor Author

Can you expand on what you mean? the class is defined in CUTLASS, so I don't see how you could break the dependency.
Or do you mean copy their implementation into MatX?

@cliffburdick
Copy link
Collaborator

Can you expand on what you mean? the class is defined in CUTLASS, so I don't see how you could break the dependency. Or do you mean copy their implementation into MatX?

Right, when we talked I thought we discussed how the cutlass implementation is so minimal and somewhat boilerplate what we could just copy the small pieces we needed into our own type and not rely on the external dependency.

@simonbyrne
Copy link
Contributor Author

Ah, I misunderstood. I thought you meant typealias the code.

The challenge with that is that it depends on the small matrix/vector types which are also defined in CUTLASS. I'll look to see how coupled they are, perhaps it is feasible to split them out.

If we do define our own quaternion type, would we still be able to use CUTLASS for quaternion GEMMs?

@cliffburdick
Copy link
Collaborator

Ah, I misunderstood. I thought you meant typealias the code.

The challenge with that is that it depends on the small matrix/vector types which are also defined in CUTLASS. I'll look to see how coupled they are, perhaps it is feasible to split them out.

If we do define our own quaternion type, would we still be able to use CUTLASS for quaternion GEMMs?

I think given that the customers who care about extra dependencies are not the same customers that care about quaternions at this point, what you have is good. We can reassess in the future if needed.

@cliffburdick
Copy link
Collaborator

Is this ready for review? It still says draft.

@cliffburdick
Copy link
Collaborator

/build

@simonbyrne simonbyrne marked this pull request as ready for review June 6, 2025 15:47
@simonbyrne
Copy link
Contributor Author

@cliffburdick It's functional, but probably needs more tests: what do you suggest?

There are still the open questions above, and what functionality we want (e.g. GEMM, abs, etc)

@cliffburdick
Copy link
Collaborator

@cliffburdick It's functional, but probably needs more tests: what do you suggest?

There are still the open questions above, and what functionality we want (e.g. GEMM, abs, etc)

I think we should have it work on the basic scalar tests like add, subtract, etc, and if it's not too hard also the common functions like abs, conj, etc. gemm would be a bonus, but likely much harder.

@simonbyrne
Copy link
Contributor Author

Sounds good. I'll try it out and see how I go.

Since this will be gated behind CUTLASS functionality, can we add that as a build option to the CI?

@cliffburdick
Copy link
Collaborator

Sounds good. I'll try it out and see how I go.

Since this will be gated behind CUTLASS functionality, can we add that as a build option to the CI?

yeah I'll add that

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.

2 participants