@triton.jit
def triton_mm(x_ptr, y_ptr, out_ptr, n: tl.constexpr, m: tl.constexpr, p:
    tl.constexpr, BLOCK_SIZE: tl.constexpr):
    pid0 = tl.program_id(axis=0)
    pid1 = tl.program_id(axis=1)
    x_row_start = tl.arange(0, BLOCK_SIZE)[:, None] * m
    x_col_start = tl.arange(0, BLOCK_SIZE)[None, :]
    x_col_unit = BLOCK_SIZE
    y_row_start = tl.arange(0, BLOCK_SIZE)[:, None] * p
    y_col_start = tl.arange(0, BLOCK_SIZE)[None, :]
    y_row_unit = BLOCK_SIZE * p
    o_row_start = pid0 * BLOCK_SIZE * p + tl.arange(0, BLOCK_SIZE)[:, None] * p
    o_col_start = pid1 * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)[None, :]
    o_offset = o_row_start + o_col_start
    out = tl.zeros((BLOCK_SIZE, BLOCK_SIZE), dtype=tl.float32)
    for i in range(tl.cdiv(m, BLOCK_SIZE)):
        x_block = x_row_start + x_col_start
        y_block = y_row_start + y_col_start
        x_offset = pid0 * BLOCK_SIZE * m + x_block
        x_mask = x_row_start < n * m and x_col_start < m
        y_offset = pid1 * BLOCK_SIZE + y_block
        y_mask = y_row_start < m * p and y_col_start < p
        x = tl.load(x_ptr + x_offset, mask=x_mask, other=0.0)
        y = tl.load(y_ptr + y_offset, mask=y_mask, other=0.0)
        out += tl.dot(x, y, allow_tf32=False)
        x_col_start += x_col_unit
        y_row_start += y_row_unit
    o_mask = o_row_start < n * p and o_col_start < p
    tl.store(out_ptr + o_offset, out, mask=o_mask)
