Skip to main content

/

Docs

Nightly
v24.5
Log in

Mojo function

tile_middle_unswitch_boundaries

tile_middle_unswitch_boundaries[work_fn: fn[Int, Bool](Int) capturing -> None, middle_tile_sizes: VariadicList[Int], left_tile_size: Int = 1, right_tile_size: Int = 1](left_boundary_start: Int, left_boundary_end: Int, right_boundary_start: Int, right_boundary_end: Int)

Divides 1d iteration space into three parts and tiles them with different steps.

The 1d iteration space is divided into: 1. [left_boundary_start, left_boundary_end), effected by left boundary. 2. [left_boundary_end, right_boundary_start), not effected by any boundary. 3. [right_boundary_start, right_boundary_end), effected by right boundary.

work_fn's switch is true for the left and right boundaries, implying boundary conditions like padding in convolution. The middle part is tiled with static tile sizes with the switch as false.

middle_tile_sizes should be in descending order for optimal performance. (Larger tile size appeared later in the list fails the while-loop.)

Parameters:

  • ​work_fn (fn[Int, Bool](Int) capturing -> None): Work function that processes one tile of workload.
  • ​middle_tile_sizes (VariadicList[Int]): List of tile sizes for the middle part.
  • ​left_tile_size (Int): Tile size for the left boundary region.
  • ​right_tile_size (Int): Tile size for the right boundary region.

Args:

  • ​left_boundary_start (Int): Start index of the left boundary.
  • ​left_boundary_end (Int): End index of the left boundary.
  • ​right_boundary_start (Int): Start index of the right boundary.
  • ​right_boundary_end (Int): End index of the right boundary.

tile_middle_unswitch_boundaries[work_fn: fn[Int, Bool, Bool](Int) capturing -> None, tile_size: Int, size: Int]()

Tile 1d iteration space with boundary conditions at both ends.

This generator is primarily for convolution with static shapes. work_fn's flags hints the function to handle padding at the boundary. The size is the static output row size, i.e., WO dimension.

Parameters:

  • ​work_fn (fn[Int, Bool, Bool](Int) capturing -> None): Work function that updates one tile. It has two flags for left and right boundaries, respectively.
  • ​tile_size (Int): 1D Tile size.
  • ​size (Int): Iteration range is [0, size).