Mojo module

swizzle

Defines swizzle layouts for optimizing memory access patterns.

This module is designed for use in shared memory, especially in GPU kernels, to reduce bank conflicts. It provides tools to create and apply swizzle transformations to memory indices. Swizzling rearranges memory access order to distribute accesses across different memory banks. This mitigates bank contention and improves memory access efficiency.

Module components:

• Swizzle struct: Represents a swizzle transformation with configurable bits, base, and shift parameters.
• Helper functions: make_ldmatrix_swizzle, make_swizzle create predefined swizzle patterns. These are optimized for scenarios like ldmatrix instructions and general 2D memory access.
• ComposedLayout struct: Combines a base layout with a swizzle layout for complex memory access optimizations.

Primary use case: GPU kernel development where shared memory bank conflicts can degrade performance. Applying swizzle layouts optimizes memory access patterns for higher throughput.

Structs

• ​ComposedLayout: Layout composed of two layouts applied sequentially.
• ​Swizzle: Swizzle functor for memory access pattern optimization.

Functions

• ​eval_composed: Evaluate a composed layout with swizzle.
• ​make_ldmatrix_swizzle: Make swizzle to avoid bank conflict for ldmatrix ops.
• ​make_swizzle: Create a 2D swizzle to avoid bank conflicts.
• ​shiftl: Shift left or right based on sign of shift amount.
• ​shiftr: Shift right or left based on sign of shift amount.

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Some code doesn’t work

It includes inaccurate information

It's missing information I need

It was difficult to understand

Other

Submit