error message improvement for hl.tile([x])

[shunting314]

For this kernel

@helion.kernel(use_default_config=True)
def helion_sum(x):
    m, n = x.shape
    y = torch.empty(m, device=x.device, dtype=x.dtype)

    for tile in hl.tile([m]):
        y[tile] = x[tile, :].sum(dim=-1)
    return y

Helion reports the following error:

  File "/home/shunting/ws/helion/helion/_compiler/type_propagation.py", line 481, in _device_indexing_size
    raise exc.OverpackedTile(k)
helion.exc.OverpackedTile: Got a tile wrapped inside a container when indexing a tensor: SequenceType((TileIndexType(0), ))
Did you mix up `hl.tile([x])` and `hl.tile(x)`?
While processing:
  File "/home/shunting/gws/tutor/helion/sum.py", line 11, in helion_sum
    y[tile] = x[tile, :].sum(dim=-1)
              ^^^^^^^^^^

The error message is clear about how to fix. But it's not clear why hl.tile should not talke [m] as argument while a tensor.shape is a valid input in 2D case.

[jansel]

You can do:

for (tile,) in hl.tile([m]):

or

for tile in hl.tile(m):

If you pass a list/tuple it returns a list/tuple.

Any suggestions to make that more clear?

[shunting314]

@jansel Got it. But without looking into how Helion work internally (from the perspective of a new user), here is why it feels a bit confusing.

@helion.kernel(use_default_config=True)
def helion_add(x, y):
    z = torch.empty_like(x)
    for tile in hl.tile(z.shape):
        z[tile] = x[tile] + y[tile]

    return z

For pointwise add, I can pass a tensor shape to hl.tile directly. While for a sum kernel (a slightly changed version):

@helion.kernel(use_default_config=True)
def helion_sum(x):
    m, n = x.shape
    y = torch.empty(m, device=x.device, dtype=x.dtype)

    for tile in hl.tile(y.shape):
        y[tile] = x[tile, :].sum(dim=-1)
    return y

it will fail.

I think it's more clear to make the second kernel work as well unless there is some fundamental reason it should not work?

[shunting314]

i.e. does it make sense for syntax like x[tile, :], flatten tile so it can be either a single tile, or a tuple of tiles.

[jansel]

In python x[(y,z)] is the same as x[y,z], so this is an artifact of that. We could make both work though.