Skip to content

Implement Swizzle Layout to replace CUTLASS Swizzle. #40

@KuangjuX

Description

@KuangjuX

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.

Metadata

Metadata

Assignees

Labels

enhancementNew feature or request

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions