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

Implement Swizzle Layout to replace CUTLASS Swizzle. #40

Open
KuangjuX opened this issue Jan 16, 2025 · 0 comments · May be fixed by #47
Open

Implement Swizzle Layout to replace CUTLASS Swizzle. #40

KuangjuX opened this issue Jan 16, 2025 · 0 comments · May be fixed by #47
Assignees
Labels
enhancement New feature or request

Comments

@KuangjuX
Copy link
Collaborator

KuangjuX commented Jan 16, 2025

To achieve the Swizzle Layout and replace the implementation of Swizzle in CUTLASS, the following functionalities need to be implemented:

  • Implement a Swizzle functor and realize the apply function, with the template as shown below:
template <const int kB, const int kM, const int kS>
struct Swizzle {
    DEVICE int32 apply(int32 idx);
};

The apply method transforms a 1D coordinate into another 1D coordinate, with the coordinate's range lying between $\left[ 0, 2^B \times 2^S \times 2^M - 1 \right]$.

In Swizzle, the index representation can be as follows:

|  XXX   | YYY   | ZZZ   |
| Bbits  | Sbits | MBits |

where the data in Mbits keep constant.

  • Implement a Swizzle Layout to replace the CUTLASS Swizzle, for example:
    using SwizzledBaseTile = decltype(composition(
        cute::Swizzle<kB, kM, kS>{},
        cute::Layout<Shape<Int<BaseShape::kRows>, Int<BaseShape::kCols>>,
                     Stride<Int<BaseShape::kCols>, _1>>{}));

in: https://github.com/microsoft/TileFusion/blob/master/include/types/layout.hpp#L40

In order to implement a SwizzleLayout, it requires the use of a Swizzle<B, M, S> along with a Layout as template parameters:

template SwizzleLayout<typename Swizzle<B, M, S>, typename Layout>{};

The amount of data in the Layout should be equal to $2^B * 2^S * 2^M$. SwizzleLayout needs to provide a method that, when accessing a coordinate within a Layout, converts it into the swizzled coordinate.

@KuangjuX KuangjuX added the enhancement New feature or request label Jan 16, 2025
@KuangjuX KuangjuX changed the title Implement the Swizzle Layout to replace CuTe Implement Swizzle Layout to replace CuTe Jan 16, 2025
@KuangjuX KuangjuX changed the title Implement Swizzle Layout to replace CuTe Implement Swizzle Layout to replace CUTLASS Swizzle. Jan 17, 2025
@KuangjuX KuangjuX self-assigned this Jan 18, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

Successfully merging a pull request may close this issue.

1 participant