Helion errors with size 0 tensors

[Chillee]

TitanTrainingActor(rank=3, ip_address=10.133.155.52) pid=3182556)   File "/tmp/horace/uv/install/6fbae45f965d50cba20076240d49c0af606f091cbea5976d1cc16bb095b35930/lib/python3.12/site-packages/helion/runtime/kernel.py", line 292, in __call__
(TorchTitanTrainingActor(rank=3, ip_address=10.133.155.52) pid=3182556)     return self.bind(args)(*args)
(TorchTitanTrainingActor(rank=3, ip_address=10.133.155.52) pid=3182556)            ^^^^^^^^^^^^^^^
(TorchTitanTrainingActor(rank=3, ip_address=10.133.155.52) pid=3182556)   File "/tmp/horace/uv/install/6fbae45f965d50cba20076240d49c0af606f091cbea5976d1cc16bb095b35930/lib/python3.12/site-packages/helion/runtime/kernel.py", line 173, in bind
(TorchTitanTrainingActor(rank=3, ip_address=10.133.155.52) pid=3182556)     bound_kernel = BoundKernel(self, args)
(TorchTitanTrainingActor(rank=3, ip_address=10.133.155.52) pid=3182556)                    ^^^^^^^^^^^^^^^^^^^^^^^
(TorchTitanTrainingActor(rank=3, ip_address=10.133.155.52) pid=3182556)   File "/tmp/horace/uv/install/6fbae45f965d50cba20076240d49c0af606f091cbea5976d1cc16bb095b35930/lib/python3.12/site-packages/helion/runtime/kernel.py", line 358, in __init__
(TorchTitanTrainingActor(rank=3, ip_address=10.133.155.52) pid=3182556)     self.host_function: HostFunction = HostFunction(
(TorchTitanTrainingActor(rank=3, ip_address=10.133.155.52) pid=3182556)                                        ^^^^^^^^^^^^^
(TorchTitanTrainingActor(rank=3, ip_address=10.133.155.52) pid=3182556)   File "/tmp/horace/uv/install/6fbae45f965d50cba20076240d49c0af606f091cbea5976d1cc16bb095b35930/lib/python3.12/site-packages/helion/_compiler/host_function.py", line 114, in __init__
(TorchTitanTrainingActor(rank=3, ip_address=10.133.155.52) pid=3182556)     propagate_types(self)
(TorchTitanTrainingActor(rank=3, ip_address=10.133.155.52) pid=3182556)   File "/tmp/horace/uv/install/6fbae45f965d50cba20076240d49c0af606f091cbea5976d1cc16bb095b35930/lib/python3.12/site-packages/helion/_compiler/type_propagation.py", line 2375, in propagate_types
(TorchTitanTrainingActor(rank=3, ip_address=10.133.155.52) pid=3182556)     prop.visit(stmt)
(TorchTitanTrainingActor(rank=3, ip_address=10.133.155.52) pid=3182556)   File "/tmp/horace/uv/install/6fbae45f965d50cba20076240d49c0af606f091cbea5976d1cc16bb095b35930/lib/python3.12/site-packages/helion/_compiler/type_propagation.py", line 1630, in visit
(TorchTitanTrainingActor(rank=3, ip_address=10.133.155.52) pid=3182556)     type_info = visitor(node)
(TorchTitanTrainingActor(rank=3, ip_address=10.133.155.52) pid=3182556)                 ^^^^^^^^^^^^^
(TorchTitanTrainingActor(rank=3, ip_address=10.133.155.52) pid=3182556)   File "/tmp/horace/uv/install/6fbae45f965d50cba20076240d49c0af606f091cbea5976d1cc16bb095b35930/lib/python3.12/site-packages/helion/_compiler/type_propagation.py", line 2179, in visit_For
(TorchTitanTrainingActor(rank=3, ip_address=10.133.155.52) pid=3182556)     iter_type = self.visit(node.iter)
(TorchTitanTrainingActor(rank=3, ip_address=10.133.155.52) pid=3182556)                 ^^^^^^^^^^^^^^^^^^^^^
(TorchTitanTrainingActor(rank=3, ip_address=10.133.155.52) pid=3182556)   File "/tmp/horace/uv/install/6fbae45f965d50cba20076240d49c0af606f091cbea5976d1cc16bb095b35930/lib/python3.12/site-packages/helion/_compiler/type_propagation.py", line 1638, in visit
(TorchTitanTrainingActor(rank=3, ip_address=10.133.155.52) pid=3182556)     raise exc.InternalError(e) from e

[oulgen]

@Chillee can you give a repro as well?

[Chillee]

@oulgen just running the add tutorial with a size 0 tensor gives this error.

import torch

import helion
from helion._testing import DEVICE
from helion._testing import run_example
import helion.language as hl

# %%
# Addition Kernel
# ---------------


# %%
@helion.kernel()
def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """
    Add two tensors element-wise with broadcasting support.

    Args:
        x: First input tensor
        y: Second input tensor

    Returns:
        A new tensor containing the element-wise sum of x and y
    """
    # match pytorch broadcasting rules
    x, y = torch.broadcast_tensors(x, y)
    out = torch.empty(
        x.shape,
        # match type promotion of torch.add
        dtype=torch.promote_types(x.dtype, y.dtype),
        device=x.device,
    )
    # tile will be a tuple of blocks
    for tile in hl.tile(out.size()):
        out[tile] = x[tile] + y[tile]
    return out


# %%
# Verification Function
# ---------------------


# %%
def check(m: int, n: int) -> None:
    """
    Verify the add kernel implementation against PyTorch's native add function.

    Args:
        m: First dimension of the test tensors
        n: Second dimension of the test tensors
    """
    x = torch.randn([m, n], device=DEVICE, dtype=torch.float16)
    y = torch.randn([m, n], device=DEVICE, dtype=torch.float16)
    run_example(add, torch.add, (x, y))


# %%
# Main Function
# -------------


# %%
def main() -> None:
    """
    Main entry point that runs the add kernel verification with 1024x1024 tensors.
    """
    check(0, 10240)


if __name__ == "__main__":
    main()

[oulgen]

@Chillee what do you think should be the expected behavior for 0dim tensors though? convert them to scalar?

[jansel]

size=0 tiles just become no-ops. Should be easy to fix.

[Chillee]

Does this require specialization on size 0 tensors?

[jansel]

I suspect if you bind a kernel with sizes >1 it will work for 0 and 1, but currently bind specializes and will recompile.