Layout Functions

Complete reference for all layout creation and transformation functions.

All layout functions operate on enigma.Layout objects. Layouts are (shape, stride) pairs that map multi-dimensional coordinates to linear memory offsets.

See Layout Algebra for a conceptual introduction.

Layout type

class Layout:
    shape:  tuple
    stride: tuple
    def __call__(self, coord) -> int: ...

---

Creating layouts

enigma.Layout(shape, stride)

Explicit constructor:

L = enigma.Layout((4, 8), (1, 4))
# L((r, c)) = r*1 + c*4

enigma.make_layout(shape, stride=None)

Alias for Layout. If stride is omitted, defaults to row-major:

L = enigma.make_layout((4, 8))          # row-major: stride=(8, 1)
L = enigma.make_layout((4, 8), (1, 4))  # column-major

enigma.make_ordered_layout(shape, order)

Creates a layout with strides determined by dimension ordering. order[0] is the fastest-varying dimension:

# 4×64, columns vary fastest (standard row-major)
L = enigma.make_ordered_layout((4, 64), order=(1, 0))

# 4×64, rows vary fastest (column-major)
L = enigma.make_ordered_layout((4, 64), order=(0, 1))

enigma.make_identity_layout(shape)

Column-major layout (mode 0 varies fastest):

L = enigma.make_identity_layout((4, 8))
# stride = (1, 4)

---

Querying layouts

enigma.size(layout, mode=None)

Total number of logical elements:

enigma.size(L)              # total: shape[0] * shape[1] * ...
enigma.size(L, mode=[0])    # size of mode 0
enigma.size(L, mode=[1])    # size of mode 1
enigma.size(tv, mode=[0])   # number of threads (in a TV layout)
enigma.size(tv, mode=[1])   # values per thread (in a TV layout)

---

Transforming layouts

enigma.coalesce(layout)

Merges adjacent modes with compatible strides into a single flat mode:

L = enigma.Layout(((4, 8),), ((1, 4),))
L2 = enigma.coalesce(L)
# L2 = Layout((32,), (1,))

enigma.complement(layout, size=None)

Returns the layout covering the elements not covered by layout within size total elements:

Lc = enigma.complement(L)

enigma.composition(a, b)

Compose a and b: treats b as a re-indexing of a:

composed = enigma.composition(outer_layout, inner_layout)

enigma.logical_divide(layout, tiler)

Divide a layout into (tile, rest). Returns a layout where mode 0 is the tile and mode 1 is the outer count:

Lt = enigma.logical_divide(L, tiler=(16, 8))

enigma.zipped_divide(layout, tiler)

Like logical_divide but applies division per mode and zips the result:

Lt = enigma.zipped_divide(L, (16, 8))
# Result shape: ((tile_m, rest_m), (tile_n, rest_n))

enigma.blocked_product(a, b)

Compute the blocked outer product of two layouts:

Lb = enigma.blocked_product(layout_a, layout_b)

enigma.recast_layout(new_bits, old_bits, layout)

Rescale strides for a different element bit width:

# View a float32 layout as float16
L_f16 = enigma.recast_layout(new_bits=16, old_bits=32, layout=L_f32)

---

Thread-value layout

enigma.make_layout_tv(thread_layout, value_layout)

The central tiling primitive. Returns a (tiler, tv_layout) pair:

thr = enigma.make_ordered_layout((4, 64), order=(1, 0))
val = enigma.make_ordered_layout((4, 4), order=(1, 0))
tiler_mn, tv_layout = enigma.make_layout_tv(thr, val)

• tiler_mn — tile shape at each level
• tv_layout — maps (thread_id, value_id) → tile coordinate

---

Tensor operations

These are used inside @enigma.jit bodies:

enigma.tensor_zipped_divide(tensor, tiler)

Partition a Tensor into tiles:

gA = enigma.tensor_zipped_divide(mA, tiler_mn)
blkA = gA[((None, None), block_idx)]

enigma.tensor_composition(tensor, tv_layout, tiler)

Map thread-value indices to a tensor fragment:

thrA = enigma.tensor_composition(blkA, tv_layout, tiler_mn)
thread_frag = thrA[(thread_idx, None)]

tensor.load() / tensor.store(value)

Vectorized load/store of a thread fragment:

vals = thread_frag.load()
out_frag.store(vals)

---

Tuple utilities

Function	Description
enigma.product(t)	Product of all elements in a (possibly nested) tuple
enigma.repeat_like(ref, val)	Create a tuple matching the structure of ref, filled with val
enigma.select(t, idx)	Select elements from a nested tuple by index