Mojo function

blocked_product

blocked_product(var layout_a: Layout, var layout_b: Layout, coalesce_output: Bool = False) -> Layout

Creates a blocked layout by combining two layouts.

This function creates a hierarchical blocked layout by combining an inner (block) and an outer (base) layout. The result is a layout where each element of the outer layout is replaced by a block defined by the inner layout.

This is particularly useful for creating tiled layouts for efficient cache utilization in tensor operations like matrix multiplication.

Example:

from layout import Layout
from layout.layout import blocked_product

# Create a 2x3 matrix layout
var matrix = Layout.row_major(2, 3)
# Define 2x2 blocks
var block = Layout.row_major(2, 2)
# Create a blocked layout with 2x2 blocks
var blocked = blocked_product(block^, matrix^)

Output:

(((2, 2), (2, 3)):((2, 12), (1, 4)))
      0    1    2    3    4    5
   +----+----+----+----+----+----+
0  |  0 |  1 |  4 |  5 |  8 |  9 |
   +----+----+----+----+----+----+
1  |  2 |  3 |  6 |  7 | 10 | 11 |
   +----+----+----+----+----+----+
2  | 12 | 13 | 16 | 17 | 20 | 21 |
   +----+----+----+----+----+----+
3  | 14 | 15 | 18 | 19 | 22 | 23 |
   +----+----+----+----+----+----+

Args:

layout_a (Layout): Inner layout. The layout for an individual block, or tile.
layout_b (Layout): Outer layout. The layout of the tiles in the output layout.
coalesce_output (Bool): Whether to coalesce the output layout. Default is False.

Returns:

Layout: A new layout representing the blocked structure