@triton.autotune(configs=autotune_configs, key=['row_dense_dim',
    'row_sparse_dim', 'col_dim'])
@triton.jit
def input_row_sparse_matmul_kernel(lhs_ptr, rhs_ptr, out_ptr,
    expert_ends_ptr, expert_pad_begins_ptr, row_dense_dim: tl.constexpr,
    row_sparse_dim: tl.constexpr, col_dim: tl.constexpr, inner_dim: tl.
    constexpr, lhs_stride_row: tl.constexpr, lhs_stride_inner: tl.constexpr,
    rhs_stride_inner: tl.constexpr, rhs_stride_col: tl.constexpr,
    out_stride_row: tl.constexpr, out_stride_col: tl.constexpr, accumulate:
    tl.constexpr, block_size_row: tl.constexpr, block_size_col: tl.
    constexpr, block_size_inner: tl.constexpr, group_size_row: tl.constexpr):
    tl.static_assert(row_sparse_dim % block_size_row == 0)
    tl.static_assert(col_dim % block_size_col == 0)
    tl.static_assert(inner_dim % block_size_inner == 0)
    tl.static_assert(row_dense_dim % row_sparse_dim == 0)
    pid_row, pid_col = tl.swizzle2d(tl.program_id(axis=0), tl.program_id(
        axis=1), row_dense_dim // block_size_row, col_dim // block_size_col,
        group_size_row)
    row_dense_offset = pid_row * block_size_row
    sparse_index = row_dense_offset // row_sparse_dim
    row_sparse_offset = row_dense_offset % row_sparse_dim
    col_offset = pid_col * block_size_col
    inner_begin = tl.load(expert_ends_ptr + sparse_index - 1, mask=
        sparse_index > 0, other=0)
    inner_end = tl.load(expert_pad_begins_ptr + sparse_index)
    inner_offset = inner_begin // block_size_inner * block_size_inner
    row_range = tl.arange(0, block_size_row)[:, None]
    col_range = tl.arange(0, block_size_col)[None, :]
    inner_range = tl.arange(0, block_size_inner) + inner_offset
    lhs_ptr += (row_sparse_offset + row_range) * lhs_stride_row
    rhs_ptr += (col_offset + col_range) * rhs_stride_col
    out_ptr += (row_dense_offset + row_range) * out_stride_row + (col_offset +
        col_range) * out_stride_col
    mask = (inner_begin <= inner_range) & (inner_range < inner_end)
    out = tl.dot(tl.load(lhs_ptr + inner_range[None, :] * lhs_stride_inner,
        mask=mask[None, :], other=0), tl.load(rhs_ptr + inner_range[:, None
        ] * rhs_stride_inner, mask=mask[:, None], other=0))
    for i in range(1, tl.cdiv(inner_end - inner_offset, block_size_inner)):
        inner_range += block_size_inner
        mask = (inner_begin <= inner_range) & (inner_range < inner_end)
        out += tl.dot(tl.load(lhs_ptr + inner_range[None, :] *
            lhs_stride_inner, mask=mask[None, :], other=0), tl.load(rhs_ptr +
            inner_range[:, None] * rhs_stride_inner, mask=mask[:, None],
            other=0))
    if accumulate:
        out += tl.load(out_ptr)
    tl.store(out_ptr, out)
