diff --git a/media/docs/cpp/cute/02_layout_algebra.md b/media/docs/cpp/cute/02_layout_algebra.md index 03578313..1b81e72f 100644 --- a/media/docs/cpp/cute/02_layout_algebra.md +++ b/media/docs/cpp/cute/02_layout_algebra.md @@ -133,11 +133,15 @@ First, a few observations: With the above, we can assume without loss of generality that `B = s:d` is a layout with integral shape and stride. We can also assume that `A` is a flattened, coalesced layout. -When `A` is integral, `A = a:b`, the result is rather trivial: `R = A o B = a:b o s:d = s:(b*d)`. But when `A` is multimodal, we need to be more careful. +When `A` is integral, `A = a:b`, the result is rather trivial: `R = A o B = a:b o s:d = s:(b*d)`. Here, the result of the composition `R` are the first `s` elements of `A` strided by `d`. -Put into words, `A o B = A o s:d`, for integral `s` and `d` means that we want (1) "remove" the first `d` elements from `A`, and then (2) "keep" the first `s` of those strided elements. +When `A` is multimodal, we need to be more careful. Put into words, `A o B = A o s:d`, for integral `s` and `d` means that we want to: -1. Removing the first `d` elements of `A` can be computed by progressively "dividing out" the first `d` elements from the shape of `A` starting from the left. For example, +1. Determine a layout that produces every `d`th element of `A`. + +The shape of this intermediate layout can be calculated by progressively "dividing out" the first `d` elements from the shape of `A` starting from the left. + +For example, * `(6,2) / 2 => (3,2)` * `(6,2) / 3 => (2,2)` * `(6,2) / 6 => (1,2)` @@ -147,9 +151,14 @@ Put into words, `A o B = A o s:d`, for integral `s` and `d` means that we want ( * `(3,6,2,8) / 9 => (1,2,2,8)` * `(3,6,2,8) / 72 => (1,1,1,4)` +To compute the strides of the strided layout, the residues of the above operation are used to scale the strides of `A`. For instance, the last example `(3,6,2,8):(w,x,y,z) / 72` with strides `(w,x,y,z)` produces `(3*w,6*x,2*x,2*z)` as the strides of the strided layout. + + As you may have noticed, we can only divide shapes by certain values and get a sensible result. This is called the **stride divisibility condition** and is statically checked in CuTe when possible. -2. Keeping the first `s` elements of the strided `A` layout can be computed by "modding out" the first `s` elements from the shape of `A` starting from the left. For example, +2. Keep the first `s` elements of the newly strided `A` so that the result has a compatible shape with `B`. This can be computed by "modding out" the first `s` elements from the shape of `A` starting from the left. + +For example, * `(6,2) % 2 => (2,1)` * `(6,2) % 3 => (3,1)` * `(6,2) % 6 => (6,1)` @@ -159,11 +168,63 @@ As you may have noticed, we can only divide shapes by certain values and get a s * `(1,2,2,8) % 2 => (1,2,1,1)` * `(1,2,2,8) % 16 => (1,2,2,4)` +This operation causes the result to have a shape that is compatible with `B`. + Again, this operation must satisfy a **shape divisibility condition** to yield a sensible result and is statically checked in CuTe when possible. From the above examples, we can construct the composition `(3,6,2,8):(w,x,y,z) o 16:9 = (1,2,2,4):(3*w,3*x,y,z)`. -#### Example 1 -- Reshape a layout into a matrix +--- + +#### Example 1 -- Worked Example of Calculating a Composition + +We provide a more complex example of composition, where both operand layouts are multi-modal to illustrate the concepts introduced above. + +``` +Functional composition, R := A o B +R(c) := (A o B)(c) := A(B(c)) + +Example +A = (6,2):(8,2) +B = (4,3):(3,1) + +1. Using the left-distributive and concatenation properties for layouts we write the composition as, + +R = A o B + = (6,2):(8,2) o (4,3):(3,1) + = ((6,2):(8,2) o 4:3, (6,2):(8,2) o 3:1) + +--- +1. Compute `(6,2):(8,2) o 4:3` + +- First, we compute the strided layout, + +(6,2):(8,2) / 3 = (6/3,2):(8*3,2) = (2,2):(24,2) + +- Next, we keep the shape compatible, + +(2,2):(24,2) % 4 = (2,2):(24,2) + +--- +2. Compute `(6,2):(8,2) o 3:1` + +- First, we compute the strided layout + +(6,2):(8,2) / 1 = (6,2):(8,2) + +- Next, we keep the shape compatible, + +(6,2):(8,2) % 3 = (3,1):(8,2) + +--- + +Putting this together and coalescing each mode, we obtain the result + +R = A o B + = ((2, 2), 3): ((24, 2), 8) +``` + +#### Example 2 -- Reshape a layout into a matrix `20:2 o (5,4):(4,1)`. Composition formulation. @@ -181,7 +242,7 @@ as a 5x4 matrix in a row-major order. 4. ` = (5,4):(8,2)`. Final composed layout. -#### Example 2 -- Reshape a layout into a matrix +#### Example 3 -- Reshape a layout into a matrix `(10,2):(16,4) o (5,4):(1,5)`