Submission 05: Contraction Interface and L2 Optimization
In this assignment you will build a high-level configuration interface for tensor contractions, implement an optimizer that manipulates those configurations, and use it to derive and benchmark an L2-optimized cuTile kernel.
All code should be written in src/.
Use FP16 data type for tensor inputs and outputs, accumulate in FP32. We assume row-major order for all tensors.
Task 1: Config Class
@dataclass
class Config():
data_type: DataType
prim_main: PrimType
dim_types: list[DimType]
exec_types: list[ExecType]
dim_sizes: list[int]
strides: list[list[int]]
prim_last: LastType = LastType.NONE
prim_first: FirstType = FirstType.ZERO
def from_config(config: Self, **kwargs):
return Config(
data_type=kwargs.get("data_type", config.data_type),
prim_main=kwargs.get("prim_main", config.prim_main),
dim_types=kwargs.get("dim_types", config.dim_types),
exec_types=kwargs.get("exec_types", config.exec_types),
dim_sizes=kwargs.get("dim_sizes", config.dim_sizes),
strides=kwargs.get("strides", config.strides),
prim_last=kwargs.get("prim_last", config.prim_last),
prim_first=kwargs.get("prim_first", config.prim_first)
)
def __str__(self):
return f"""Config(
data_type={self.data_type},
prim_main={self.prim_main},
prim_last={self.prim_last},
prim_first={self.prim_first},
dim_types={self.dim_types},
exec_types={self.exec_types},
dim_sizes={self.dim_sizes},
strides={self.strides}
)"""
Task 2: Generating a Basic Config
Write a function generate_config that takes an einsum string and a list of shapes for the input tensors (the output shape is implied by the einsum) and returns a basic Config.
def generate_config(einsum: str, input_shapes: list[tuple[int]], dim_order: str | None = None) -> Config:
# 3 catpure groups: 1 for the output, 2 for the inputs, ignore whitespaces
einsum = re.sub(r'\s+', '', einsum)
A_dims, B_dims, C_dims = re.match(r"([a-z]+),([a-z]+)->([a-z]+)", einsum).groups()
# to keep the order of the dimensions as they appear in the einsum (and in the lecture)
def remove_duplicates_keep_order(seq):
seen = set()
seen_add = seen.add
return [x for x in seq if not (x in seen or seen_add(x))]
if dim_order is not None:
dim_names = list(dim_order)
else:
dim_names = remove_duplicates_keep_order(A_dims + B_dims + C_dims)
dim_types = []
dim_sizes = []
for dim in dim_names:
if dim in A_dims and dim in B_dims and dim in C_dims:
dim_type = DimType.C
elif dim in A_dims and dim in C_dims and not dim in B_dims:
dim_type = DimType.M
elif dim in B_dims and dim in C_dims and not dim in A_dims:
dim_type = DimType.N
elif dim in A_dims and dim in B_dims and not dim in C_dims:
dim_type = DimType.K
else:
raise ValueError(f"Dimension {dim} does not fit into M, N, K, C categories.")
dim_types.append(dim_type)
# Determine the size of the dimension from the input shapes
if dim in A_dims:
size = input_shapes[0][A_dims.index(dim)]
elif dim in B_dims:
size = input_shapes[1][B_dims.index(dim)]
else:
raise ValueError(f"Dimension {dim} not found in any input tensor.")
dim_sizes.append(size)
exec_types = [ExecType.SEQ] * len(dim_types)
size_of = dict(zip(dim_names, dim_sizes))
strides = []
for tensor_dims in [A_dims, B_dims, C_dims]:
own_strides = {}
stride = 1
for name in tensor_dims[::-1]:
own_strides[name] = stride
stride *= size_of[name]
# Map onto the global dimension slots; 0 means "not in this tensor".
strides.append([own_strides.get(name, 0) for name in dim_names])
return Config(
data_type=DataType.FLOAT16,
prim_main=PrimType.GEMM,
dim_types=dim_types,
exec_types=exec_types,
dim_sizes=dim_sizes,
strides=strides
)
Task 3: Optimizer Class
Implement a class Optimizer that wraps a Config and exposes methods to transform it.
a) Implement the function split_dim(dim_id: int, outer_size: int, inner_size: int).
def split_dim(self, dim_id: int, outer_size: int| None = None, inner_size: int| None = None):
if dim_id < 0:
dim_id += len(self.config.dim_types)
original_size = self.config.dim_sizes[dim_id]
if outer_size is None and inner_size is None:
raise ValueError("At least one of outer_size or inner_size must be provided.")
if outer_size is None:
if original_size % inner_size != 0:
raise ValueError(f"Inner size {inner_size} does not divide original size {original_size} and outer size is not provided.")
outer_size = original_size // inner_size
if inner_size is None:
if original_size % outer_size != 0:
raise ValueError(f"Outer size {outer_size} does not divide original size {original_size} and inner size is not provided.")
inner_size = original_size // outer_size
if outer_size * inner_size != original_size:
raise ValueError(f"Outer size {outer_size} and inner size {inner_size} do not multiply to original size {original_size}")
new_config = Config(
data_type=self.config.data_type,
prim_main=self.config.prim_main,
dim_types=self.config.dim_types[:dim_id] + [self.config.dim_types[dim_id]] * 2 + self.config.dim_types[dim_id+1:],
exec_types=self.config.exec_types[:dim_id] + [self.config.exec_types[dim_id]] * 2 + self.config.exec_types[dim_id+1:],
dim_sizes=self.config.dim_sizes[:dim_id] + [outer_size, inner_size] + self.config.dim_sizes[dim_id+1:],
# strides=self.config.strides[:dim_id] + [[self.config.strides[0][dim_id] * inner_size, self.config.strides[0][dim_id]]] + self.config.strides[dim_id+1:],
strides=[stride[:dim_id] + [stride[dim_id] * inner_size, stride[dim_id]] + stride[dim_id+1:] for stride in self.config.strides],
prim_last=self.config.prim_last,
prim_first=self.config.prim_first
)
self.config = new_config
return new_config
b) Implement the function fuse_dims(dim_id_a: int, dim_id_b: int).
def fuse_dims(self, dim_id_a: int, dim_id_b: int):
# set dim_id_a to be the smaller one to simplify the logic
if dim_id_a > dim_id_b:
dim_id_a, dim_id_b = dim_id_b, dim_id_a
if self.config.dim_types[dim_id_a] != self.config.dim_types[dim_id_b]:
raise ValueError(f"Dimensions {dim_id_a} and {dim_id_b} have different types and cannot be fused.")
for i, stride in enumerate(self.config.strides):
if stride[dim_id_a] == 0: # if both have same type and one of them is 0, the other must also be 0
continue
if not (stride[dim_id_a] == stride[dim_id_b] * self.config.dim_sizes[dim_id_b] or
stride[dim_id_b] == stride[dim_id_a] * self.config.dim_sizes[dim_id_a]):
raise ValueError(f"Dimensions {dim_id_a} and {dim_id_b} are not contiguous in dim {i} and cannot be fused.")
new_dim_sizes = list(self.config.dim_sizes)
new_dim_sizes[dim_id_a] = self.config.dim_sizes[dim_id_a] * self.config.dim_sizes[dim_id_b]
del new_dim_sizes[dim_id_b]
def new_stride(stride):
new_stride = list(stride)
new_stride[dim_id_a] = stride[dim_id_b] # stride of the fused dimension is the same as the second dimension
del new_stride[dim_id_b]
return new_stride
new_config = Config(
data_type=self.config.data_type,
prim_main=self.config.prim_main,
dim_types=self.config.dim_types[:dim_id_b] + self.config.dim_types[dim_id_b+1:], # keeps dim_type of a
exec_types=self.config.exec_types[:dim_id_b] + self.config.exec_types[dim_id_b+1:], # keeps exec_type of a
dim_sizes=new_dim_sizes,
strides=[new_stride(stride) for stride in self.config.strides],
prim_last=self.config.prim_last,
prim_first=self.config.prim_first
)
self.config = new_config
return new_config
c) Implement the function permute_dims(permutation: list[int]).
def permute_dims(self, permutation: list[int]):
if sorted(permutation) != list(range(len(self.config.dim_types))):
raise ValueError(f"New order {permutation} is not a valid permutation of dimensions.")
new_config = Config(
data_type=self.config.data_type,
prim_main=self.config.prim_main,
dim_types=[self.config.dim_types[i] for i in permutation],
exec_types=[self.config.exec_types[i] for i in permutation],
dim_sizes=[self.config.dim_sizes[i] for i in permutation],
strides=[ [stride[i] for i in permutation] for stride in self.config.strides],
prim_last=self.config.prim_last,
prim_first=self.config.prim_first
)
self.config = new_config
return new_config
d) Implement the function make_executable().
def make_executable(self):
ndims = len(self.config.dim_types)
permutation = list(range(ndims))
# find right most N, M, K dimensions, permute to the right, and set them to PRIM
def find_rightmost_dim_id(dim_type: DimType) -> int:
return max((i for i, dt in enumerate(self.config.dim_types) if dt == dim_type), default=-1)
for dim_type in [DimType.M, DimType.N, DimType.K]:
dim_id = find_rightmost_dim_id(dim_type)
if dim_id == -1:
raise ValueError("Cannot make config executable because it is missing at least one of M, N, K dimensions.")
# move the dimension to the rightmost position
permutation.remove(dim_id)
permutation.append(dim_id)
# move K dimensions to the right
permutation[:-3] = sorted(permutation[:-3], key=lambda i: self.config.dim_types[i] == DimType.K)
self.permute_dims(permutation)
# set exec types: rightmost 3 dimensions to PRIM, other K dimensions to SEQ, rest to PAR
sequential_k_dims = sum(1 for dt in self.config.dim_types if dt == DimType.K) - 1
self.config.exec_types = [ExecType.PAR] * (ndims - sequential_k_dims - 3) + [ExecType.SEQ] * sequential_k_dims + [ExecType.PRIM] * 3
self.verify() # should not raise an error
return self.config
e) Implement the function verify().
def verify(self):
# No K-dimension may have exec_type = PAR.
for dim_type, exec_type in zip(self.config.dim_types, self.config.exec_types):
if dim_type == DimType.K and exec_type == ExecType.PAR:
raise ValueError("K-dimension cannot have parallel execution type.")
# order: PAR -> SEQ -> PRIM
order = {ExecType.PAR: 0, ExecType.SEQ: 1, ExecType.PRIM: 2}
sorted_exec_types = sorted(self.config.exec_types, key=lambda x: order[x])
if self.config.exec_types != sorted_exec_types:
raise ValueError("Execution types must be in order: PAR -> SEQ -> PRIM.")
# The rightmost dimension must be PRIM and the PRIM dimensions must include at least one dimension of each type M, N, and K.
prim_dim_types = set(dim_type for dim_type, exec_type in zip(self.config.dim_types, self.config.exec_types) if exec_type == ExecType.PRIM)
if not prim_dim_types.issuperset({DimType.M, DimType.N, DimType.K}):
raise ValueError("The rightmost dimensions must be PRIM and the PRIM dimensions must include at least one dimension of each type M, N, and K.")
Task 4: L2-Optimized Batched Contraction
NOTE: wenn ab 5 dimensionen performance schlechter, dann siehe assume_div_by hints der optional task in week 2.
Consider the batched matrix multiplication expressed as cmk, ckn -> cmn with dimension sizes \(|c| = 4\), \(|m| = |n| = |k| = 4096\).
a) Use your generate_config function from Task 2 to produce the initial Config for this contraction. Report the resulting config.
Output:
Config(
data_type=DataType.FLOAT16,
prim_main=PrimType.GEMM,
prim_last=LastType.NONE,
prim_first=FirstType.ZERO,
dim_types=[<DimType.C: 3>, <DimType.M: 0>, <DimType.K: 2>, <DimType.N: 1>],
exec_types=[<ExecType.SEQ: 0>, <ExecType.SEQ: 0>, <ExecType.SEQ: 0>, <ExecType.SEQ: 0>],
dim_sizes=[4, 4096, 4096, 4096],
strides=[[16777216, 4096, 1, 0], [16777216, 0, 4096, 1], [16777216, 4096, 0, 1]]
)
b) Use your Optimizer and the implemented functions from Task 3 to transform the basic config into an L2-optimized one, following the general L2-reuse pattern from the lecture.
config.dim_sizes = [ [...], |m_l2|, |n_l2|, |m_prim|, |n_prim|, |k_prim|]
Choose the sizes for m_l2, m_prim, n_l2, n_prim and justify your choice with respect to L2 cache reuse.
Report the final config.
Output:
Config(
data_type=DataType.FLOAT16,
prim_main=PrimType.GEMM,
prim_last=LastType.NONE,
prim_first=FirstType.ZERO,
dim_types=[<DimType.C: 3>, <DimType.M: 0>, <DimType.N: 1>, <DimType.M: 0>, <DimType.N: 1>, <DimType.K: 2>, <DimType.M: 0>, <DimType.N: 1>, <DimType.K: 2>],
exec_types=[<ExecType.SEQ: 0>, <ExecType.SEQ: 0>, <ExecType.SEQ: 0>, <ExecType.SEQ: 0>, <ExecType.SEQ: 0>, <ExecType.SEQ: 0>, <ExecType.SEQ: 0>, <ExecType.SEQ: 0>, <ExecType.SEQ: 0>],
dim_sizes=[4, 4, 4, 16, 16, 32, 64, 64, 128],
strides=[[16777216, 4194304, 0, 262144, 0, 128, 4096, 0, 1], [16777216, 0, 1024, 0, 64, 524288, 0, 1, 4096], [16777216, 4194304, 1024, 262144, 64, 0, 4096, 1, 0]]
)
prim sizes
First we want to choose the optimal m_prim, n_prim, k_prim sizes of one mma instruction.
We want to fit as many primitive operations of a matrix multiplication into one operation of mma as possible which is m_prim * n_prim * k_prim.
In one mma operation the maximal memory is th max shared memory per block, which is 48 KiB for our machine. Using FP32 (4 bytes) for accumulation and FP16 (2 bytes) for inputs, the required memory for one mma is mma_size = (2 * m_prim * k_prim + 2 * k_prim * n_prim + 4 * m_prim * n_prim) bytes.
Therefore want to maximize m_prim * n_prim * k_prim s.t. mma_size = max_shared_memory_per_block.
Assuming m_prim = n_prim due to symmetry, leaves the optimization problem:
This is solved for m_prim = sqrt(max_shared_memory_per_block/12)
In the case of 48 KiB of shared memory, m_prim = n_prim = 64 and k_prim = max_shared_memory_per_block / (4 * m_prim) - m_prim = 128 is optimal.
L2 sizes
Next we want to choose m_l2 and n_l2 such that a swizzle block fits into the L2 cache.
The size of the L2 cache on our machine is 24MiB.
Every kernel loads k_outer * m_prim * k_prim values from matrix A and k_outer * n_prim * k_prim from B, respectively.
The amount of loaded values from A within a swizzle block is linear in m_l2 and for B linear in n_l2.
Therefore, in total we need to fit m_l2 * k_outer * m_prim * k_prim + n_l2 * k_outer * n_prim * k_prim values into the L2 cache.
For optimal L2 use, we want to fill ~90-95% of the 24 MiB cache
Inserting the previously calculated values, we land on:
We want to maximize the values calculated per swizzle block, thus maximizing \(m_{L2} \cdot n_{L2}\). Therefore, good values would be \(m_{L2} = n_{L2} = 16\). While only using 67% of the L2 cache, we stick to sizes being a power of 2.
Implementation
c) Implement the kernel
Implement a cuTile kernel that computes cmk, ckn -> cmn following your optimized config from b). Verify correctness of your kernel.
d) Use triton.testing.do_bench (or a similar benchmark function provided by cuTile/Torch) to measure the average kernel runtime. Report the achieved performance in TFLOPS.
Compare the performance of your L2-optimized kernel to a baseline kernel that maps BIDs in plain row-major order over (c, m, n) without any splitting or permuting. Report your findings.
def task_c_and_d():
c = 4
n = m = k = 4096
# optimal config from b:
m_prim = n_prim = 64
k_prim = 128
m_l2 = 16
n_l2 = 16
k_outer = k // k_prim
m_outer = m // (m_l2 * m_prim)
n_outer = n // (n_l2 * n_prim)
# c,m_outer,n_outer,m_l2,n_l2,k_outer,m_prim,n_prim,k_prim
A = torch.randn((c,m_outer,m_l2,k_outer,m_prim,k_prim), device='cuda', dtype=torch.float16)
B = torch.randn((c,n_outer,n_l2,k_outer,n_prim,k_prim), device='cuda', dtype=torch.float16)
C = torch.empty((c,m_outer,n_outer,m_l2,n_l2,m_prim,n_prim), device='cuda', dtype=torch.float16)
grid = (c * m_outer * n_outer * m_l2 * n_l2, 1, 1)
args = (A, B, C, m_outer, n_outer, m_prim, n_prim, k_prim, k_outer, m_l2, n_l2)
ms = triton.testing.do_bench(lambda: ct.launch(torch.cuda.current_stream(), grid, multiply, args))
tflops = 2 * (n * m * k * c) / (ms / 1000) / (10**12)
print(f"Execution time of optimized kernel (6D tensors): {ms:.2f} ms")
print(f"TFLOPS of optimized kernel (6D tensors): {tflops:.2f}")
# permute to original shape
A = A.permute(0, 1, 2, 4, 3, 5).reshape((c,m,k))
B = B.permute(0, 3, 5, 1, 2, 4).reshape((c,k,n))
C = C.permute(0, 1, 3, 5, 2, 4, 6).reshape((c,m,n))
expected = torch.einsum("cmk, ckn -> cmn", A, B)
assert torch.allclose(C, expected, atol=1e-1), "The result of c) is incorrect!"
# swizzling on 3D tensors
C = torch.empty((c,m,n), device='cuda', dtype=torch.float16)
args_3d = (A, B, C, m_outer, n_outer, m_prim, n_prim, k_prim, k_outer, m_l2, n_l2)
ms_3d = triton.testing.do_bench(lambda: ct.launch(torch.cuda.current_stream(), grid, multiply_3d, args_3d))
assert torch.allclose(C, expected, atol=1e-1), "The result of the 3D kernel is incorrect!"
tflops_3d = 2 * (n * m * k * c) / (ms_3d / 1000) / (10**12)
print(f"Execution time of optimized kernel (3D tensors): {ms_3d:.2f} ms")
print(f"TFLOPS of optimized kernel (3D tensors): {tflops_3d:.2f}")
# no swizzling
C = torch.empty((c,m,n), device='cuda', dtype=torch.float16)
args_baseline = (A, B, C, m_prim, n_prim, k_prim, k//k_prim)
grid_baseline = (c, m//m_prim, n//n_prim)
ms_baseline = triton.testing.do_bench(lambda: ct.launch(torch.cuda.current_stream(), grid_baseline, baseline_multiply, args_baseline))
assert torch.allclose(C, expected, atol=1e-1), "The result of baseline is incorrect!"
tflops_baseline = 2 * (n * m * k * c) / (ms_baseline / 1000) / (10**12)
print(f"Execution time of baseline kernel: {ms_baseline:.2f} ms")
print(f"TFLOPS of baseline kernel: {tflops_baseline:.2f}")
with open(file_dir / 'task4_results.txt', 'w') as f:
f.write(f"Execution time of optimized kernel (6D tensors): {ms:.2f} ms\n")
f.write(f"TFLOPS of optimized kernel (6D tensors): {tflops:.2f}\n")
f.write(f"Execution time of optimized kernel (3D tensors): {ms_3d:.2f} ms\n")
f.write(f"TFLOPS of optimized kernel (3D tensors): {tflops_3d:.2f}\n")
f.write(f"Execution time of baseline kernel: {ms_baseline:.2f} ms\n")
f.write(f"TFLOPS of baseline kernel: {tflops_baseline:.2f}\n")
plot_results(tflops, tflops_3d, tflops_baseline)
def swizzle_position(pid, m_outer, n_outer, m_l2, n_l2):
l2_group_size = m_l2 * n_l2
outer_group_size = m_outer * n_outer
c_it = pid // (outer_group_size * l2_group_size)
rem = pid % (outer_group_size * l2_group_size)
mn_outer_it = rem // l2_group_size
mn_l2_it = rem % l2_group_size
m_outer_it = mn_outer_it // n_outer
n_outer_it = mn_outer_it % n_outer
m_l2_it = mn_l2_it // n_l2
n_l2_it = mn_l2_it % n_l2
return c_it, m_outer_it, n_outer_it, m_l2_it, n_l2_it
@ct.kernel
def multiply(A, B, C, m_outer: ct.Constant[int], n_outer: ct.Constant[int], m_prim: ct.Constant[int], n_prim: ct.Constant[int], k_prim: ct.Constant[int], k_outer: ct.Constant[int], m_l2: ct.Constant[int], n_l2: ct.Constant[int]):
pid = ct.bid(0)
c_it, m_outer_it, n_outer_it, m_l2_it, n_l2_it = swizzle_position(pid, m_outer, n_outer, m_l2, n_l2)
acc = ct.zeros((m_prim, n_prim), dtype=ct.float32)
for k_it in range(k_outer):
# c,m_outer,n_outer,m_l2,n_l2,k_outer,m_prim,n_prim,k_prim
A_tile = ct.load(
A,
index=(c_it,m_outer_it,m_l2_it,k_it,0,0),
shape=(1,1,1,1,m_prim,k_prim),
).reshape((m_prim, k_prim))
B_tile = ct.load(
B,
index=(c_it,n_outer_it,n_l2_it,k_it,0,0),
shape=(1,1,1,1,n_prim,k_prim),
).reshape((n_prim,k_prim)).transpose()
acc = ct.mma(A_tile, B_tile, acc=acc)
C_ = acc.astype(ct.float16).reshape((1,1,1,1,1,m_prim,n_prim))
ct.store(C, index=(c_it,m_outer_it,n_outer_it,m_l2_it,n_l2_it,0,0), tile=C_)
@ct.kernel
def multiply_3d(A, B, C, m_outer: ct.Constant[int], n_outer: ct.Constant[int], m_prim: ct.Constant[int], n_prim: ct.Constant[int], k_prim: ct.Constant[int], k_outer: ct.Constant[int], m_l2: ct.Constant[int], n_l2: ct.Constant[int]):
pid = ct.bid(0)
c_it, m_outer_it, n_outer_it, m_l2_it, n_l2_it = swizzle_position(pid, m_outer, n_outer, m_l2, n_l2)
m_it = m_outer_it * m_l2 + m_l2_it
n_it = n_outer_it * n_l2 + n_l2_it
acc = ct.zeros((1, m_prim, n_prim), dtype=ct.float32)
for k_it in range(k_outer):
A_val = ct.load(A, index=(c_it, m_it, k_it), shape=(1, m_prim, k_prim))
B_val = ct.load(B, index=(c_it, k_it, n_it), shape=(1, k_prim, n_prim))
acc = ct.mma(A_val, B_val, acc=acc)
C_ = acc.astype(ct.float16)
ct.store(C, index=(c_it, m_it, n_it), tile=C_)
Results
multiply was our first implementation. It reshapes A, B, C into 6D tensors mirroring the tiling hierarchy (c, m_outer, m_l2, k_outer, m_prim, k_prim, …), so each output tile is reached by indexing directly into its (m_outer, m_l2) / (n_outer, n_l2) position, matching the config from b) closely. It ended up slower than the row-major baseline.
To find out why, we added multiply_3d. It uses the same swizzle_position block-id decomposition, so both kernels compute the same (m, n) tile for a given block in the same order. The only difference is that multiply_3d indexes directly into the (c, m, k) / (c, k, n) / (c, m, n) tensors instead of the reshaped 6D ones. That separates two possible explanations: a bad swizzle pattern, or overhead from the 6D tensor layout.
multiply’s 6D loads carry four extra size-1 index dimensions per ct.load, which multiply_3d avoids. Once that overhead is gone, multiply_3d beats the baseline.
Execution time of optimized kernel (6D tensors): 89.64 ms
TFLOPS of optimized kernel (6D tensors): 6.13
Execution time of optimized kernel (3D tensors): 11.31 ms
TFLOPS of optimized kernel (3D tensors): 48.59
Execution time of baseline kernel: 37.43 ms
TFLOPS of baseline kernel: 14.69
Update
after building cutile with TMA_MAX_NDIM = 10 the performance of the 6D kernel improved significantly:
Execution time of optimized kernel (6D tensors): 14.28 ms
TFLOPS of optimized kernel (6D tensors): 38.49
Execution time of optimized kernel (3D tensors): 13.04 ms
TFLOPS of optimized kernel (3D tensors): 42.15
Execution time of baseline kernel: 40.32 ms
TFLOPS of baseline kernel: 13.63
