Kernel Composition

Author: Hongzheng Chen (hzchen@cs.cornell.edu)

This document will discuss kernel composition. In the previous tutorials, we have seen how to write a simple kernel. However, in real applications, we often need to compose multiple kernels together.

In the following example, we define a matrix_add and a gemm kernel, and wrap them into a top-level function.

import allo
from allo.ir.types import int32, float32

M, K, N = 32, 32, 32


def matrix_add(A: int32[M, N]) -> int32[M, N]:
    B: int32[M, N] = 0
    for i, j in allo.grid(M, N):
        B[i, j] = A[i, j] + 1
    return B


def gemm(A: int32[M, K], B: int32[K, N]) -> int32[M, N]:
    C: int32[M, N] = 0
    for i, j in allo.grid(M, N):
        for k in allo.reduction(K):
            C[i, j] += A[i, k] * B[k, j]
    return C


def top(A: int32[M, K], B: int32[K, N]) -> int32[M, N]:
    C = gemm(A, B)
    D = matrix_add(C)
    return D

Different teams or people can then work on different parts of the code and optimize each kernel. We first create a schedule for the matrix_add kernel, and add several optimizations.

s1 = allo.customize(matrix_add)
s1.pipeline("j")
print(s1.module)

Then we create a schedule for the gemm kernel and optimize it.

s2 = allo.customize(gemm)
s2.reorder("k", "j")
s2.buffer_at(s2.C, axis="i")
s2.pipeline("j")
print(s2.module)

Notice that now we only optimize the separate kernels but do not incorporate them into the top-level function, as shown in the following printed module.

s = allo.customize(top)
print(s.module)

Therefore, after each part has been optimized, we need to explicitly compose them together. In Allo, we can use the .compose() primitive to compose the schedules together into the parent function.

s.compose([s1, s2])
print(s.module)

We can see that the schedules for the matrix_add and gemm kernels are both correctly optimized in the top-level function.

Template Composition

Sometimes we may define template kernels and invoke the kernel with different template arguments. Allo provides an id option to specify the exact kernel to be composed.

def kernel[T_in, T_out, S](A: "T_in[S]") -> "T_out[S]":
    B: T_out[S] = 0
    for i in range(S):
        with allo.meta_if(T_out == int32):
            B[i] = A[i] + 1
        with allo.meta_else():
            B[i] = A[i] * 2
    return B


def top2(A: int32[M]) -> float32[M]:
    C = kernel[int32, int32, M, "K1"](A)
    D = kernel[int32, float32, M, "K2"](C)
    return D

Specifically, the last argument of the template kernel is the id of the kernel. Later on we can use this ID for distinguishing different kernels during composition. We also customize the two template kernels with different optimizations first.

s1 = allo.customize(kernel, instantiate=[int32, int32, M])
s1.unroll("i", factor=4)
print(s1.module)

s2 = allo.customize(kernel, instantiate=[int32, float32, M])
s2.pipeline("i")
print(s2.module)

Finally, we compose the two template kernels into the top-level function with the ID specified.

s = allo.customize(top2)
s.compose(s1, id="K1")
s.compose(s2, id="K2")
print(s.module)

We can see from the printed module that the loop in the first kernel is unrolled by a factor of 4, and the loop in the second kernel is pipelined.