Pad to next power of 2 for hl.specialize'ed shape value used in device tensor creation

[yf225]

We want to support usage pattern like:

@helion.kernel()
def reduce_kernel(x: torch.Tensor) -> torch.Tensor:
    m_block = hl.register_block_size(x.size(0))
    grad_weight = x.new_empty(
        [(x.size(0) + m_block - 1) // m_block, x.size(1)], dtype=torch.float32
    )
    weight_shape = hl.specialize(x.size(1))
    for mb_cta in hl.tile(x.size(0), block_size=m_block):
        grad_w_m = x.new_zeros(weight_shape, dtype=torch.float32)
        for mb in hl.tile(mb_cta.begin, mb_cta.end):
            grad_w_m += x[mb, :].to(torch.float32).sum(0)
        grad_weight[mb_cta.id, :] = grad_w_m
    return grad_weight.sum(0).to(x.dtype)

Particularly grad_w_m += x[mb, :].to(torch.float32).sum(0) is difficult to support, because the LHS grad_w_m is of shape x.size(1) which can be a non-power-of-2 value (e.g. 56), while the RHS x[mb, :].to(torch.float32).sum(0) is a the next power-of-2 value of x.size(1) (i.e. 64), resulting in a shape mismatch.

I explored many solutions, and the cleanest / simplest way is to bump n to next power-of-2 when it's used as device tensor shape value (e.g. when used in torch.zeros in device loop).

Alternative is we ask users to explicitly wrap weight_shape with helion.next_power_of_2(...), like grad_w_m = x.new_zeros(helion.next_power_of_2(weight_shape), dtype=torch.float32), but I feel that the UX friction is quite high if we do that.

Fixes #737.
Fixes #741.

cc. @v0i0 @mengluy0125

[jansel]

I think it would be better to do this at the aten level rather than monkey-patching torch.

At the aten level the factory functions are normalized (so you need to handle fewer of them) and if you have multiple Python threads then monkey patching torch is not thread safe.

[yf225]

Updated to use TorchDispatchMode to intercept tensor factory ops at the aten level and avoid monkey-patching.