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.
To achieve the
Swizzle Layoutand replace the implementation ofSwizzlein CUTLASS, the following functionalities need to be implemented:The$\left[ 0, 2^B \times 2^S \times 2^M - 1 \right]$ .
applymethod transforms a 1D coordinate into another 1D coordinate, with the coordinate's range lying betweenIn Swizzle, the index representation can be as follows:
where the data in
Mbitskeep constant.in: https://github.com/microsoft/TileFusion/blob/master/include/types/layout.hpp#L40
In order to implement a
SwizzleLayout, it requires the use of aSwizzle<B, M, S> along with aLayoutas template parameters:The amount of data in the Layout should be equal to$2^B * 2^S * 2^M$ .
SwizzleLayoutneeds to provide a method that, when accessing a coordinate within a Layout, converts it into the swizzled coordinate.