Add hl.register_block_size and explicit tile sizes

[jansel]

This change makes the two pass softmax @drisspg was trying to implement possible, with some new syntax for pre-registering block sizes:

@helion.kernel(config={"block_sizes": [1, 128]})
def softmax_two_pass(x: torch.Tensor) -> torch.Tensor:
    m, n = x.size()
    out = torch.empty_like(x)
    block_size_m = hl.register_block_size(m)
    block_size_n = hl.register_block_size(n)
    for tile_m in hl.tile(m, block_size=block_size_m):
        mi = hl.full([tile_m, 1], float("-inf"), dtype=torch.float32)
        di = hl.zeros([tile_m, block_size_n], dtype=torch.float32)
        for tile_n in hl.tile(n, block_size=block_size_n):
            values = x[tile_m, tile_n]
            local_amax = torch.amax(values, dim=1, keepdim=True)
            mi_next = torch.maximum(mi, local_amax)
            di = di * torch.exp(mi - mi_next) + torch.exp(values - mi_next)
            mi = mi_next
        for tile_n in hl.tile(n, block_size=block_size_n):
            values = x[tile_m, tile_n]
            out[tile_m, tile_n] = torch.exp(values - mi) / di
    return out

The part that was hard before was referencing block_size_n before the tile_n loop.

I also considered something like:

di = hl.zeros([tile_m, hl.auto()], dtype=torch.float32)

Where the hl.auto() would be a placeholder value that we infer automatically during type propagation, though I think this version is clearer.