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

Add LL::quotient and remove uses of divideRight and sublayoutIsIdentity #4968

Merged
merged 7 commits into from
Oct 24, 2024

Conversation

lezcano
Copy link
Contributor

@lezcano lezcano commented Oct 22, 2024

We add a new abstraction LL::quotient that abstracts the idea of "a
linear layout does not permute certain dimensions". Doing so, allows us
to remove divideRight and subsume them into this higher-level abstraction.
We also fix a bug in isCrossCTAConversion.

We also remove some code duplication from transferWithinThreads and
cvtReorderRegisters in favour of a more generic approach.

We fix a bug in sublayout that meant that sublayout would reorder outDims
at will by using a set instead of a vector.

I am missing adding tests for LL::quotient, will do in a minute.

Copy link
Contributor

@Jokeren Jokeren left a comment

Choose a reason for hiding this comment

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

May I know what bug is there?

We also fix a bug in isCrossCTAConversion.

include/triton/Tools/LinearLayout.h Outdated Show resolved Hide resolved
include/triton/Tools/LinearLayout.h Show resolved Hide resolved
// Is the sublayout defined from dimNames to dimNames the identity?
// In particular, is the input and output size in these dimensions
// the same, and are the bases the identity?
bool squareSublayoutIsIdentity(ArrayRef<StringAttr> dimNames) const;
Copy link
Contributor

Choose a reason for hiding this comment

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

So, are we ignoring the case where the input dimension names are not equal to the output dimension names, but the input and output are identity mappings?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

The identity in maths is just defined between a space and itself, so yes, this is the correct mathematical concept. Even if you have registers that map 1-to-1 to a vector of outputs, this would not be an identity map strictly speaking, as it's not mapping elements to themselves, but identifying registers with a matrix.

Copy link
Contributor

Choose a reason for hiding this comment

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

OK, makes sense then

@lezcano
Copy link
Contributor Author

lezcano commented Oct 22, 2024

May I know what bug is there?

Previously, we were checking that:

return !layout.sublayoutIsZero(nonBlockInDims.getArrayRef(), {kBlock}) ||
         !layout.sublayoutIsIdentity({kBlock}, {kBlock});

This means that the kBlock of the output will be defined uninquely and as the identity by the kBlock of the input, which is fine. But we are missing a check that layout.sublayoutIsZero({kBlock}, nonBlockOutDims.getArrayRef()), as otherwise the layout will depend on the kBlock, but a few lines after that we are "removing" the kBlock dimensions, effectively assuming that that sublayout is zero.

At the moment I have drawn this rather rough distribution, but it can be optimised in the future. For example, if you just have those two initial conditions, you can still perform the conversion without using distributed shmem, but you'll need to write an indexing that depends on kBlock, which is a tad trickier, but of course, can be done.

@Jokeren
Copy link
Contributor

Jokeren commented Oct 22, 2024

This means that the kBlock of the output will be defined uninquely and as the identity by the kBlock of the input, which is fine. But we are missing a check that layout.sublayoutIsZero({kBlock}, nonBlockOutDims.getArrayRef()), as otherwise the layout will depend on the kBlock, but a few lines after that we are "removing" the kBlock dimensions, effectively assuming that that sublayout is zero.

Sure, I think it makes sense

We add a new abstraction `LL::quotient` that abstracts the idea of "a
linear layout does not permute certain dimensions". Doing so, allows us
to remove `divideRight` and subsume them into this higher-level abstraction.
We also fix a bug in `isCrossCTAConversion`.

We alsoremove some code duplication from `transferWithinThreads` and
`cvtReorderRegisters` in favour of a more generic approach.
@lezcano lezcano merged commit 13594bb into triton-lang:main Oct 24, 2024
7 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.

2 participants