device non grid forloops cannot subscript

[oulgen]

        @helion.kernel
        def kernel(As: list[torch.Tensor]) -> torch.Tensor:
            out = torch.zeros_like(As[0])
            for tile in hl.tile(out.size()):
                for i in range(len(As)):
                    a = As[i]
                    out[tile] += a[tile]
            return out

        args = [torch.randn(16, device=DEVICE) for _ in range(4)]
        code, result = code_and_output(kernel, (args,))
        torch.testing.assert_close(result, sum(args))

doesnt work because type_prop thinks that i is GridIndexType and there's no getitem for tensor with GridIndexType.

[oulgen]

@jansel I think this is an interesting case where the helion type system is too literal and unable to support this. Currently there's no way to say an item from the list[Tensor], we have to know exactly which item otherwise subscript As[i] fails.

One way to fix this would be to unroll the inner for loop, or even rewrite the range(len(As)) as [0..len(As)] in the literal form to avoid the GridIndexType coming from range handling.

what do you think?

[jansel]

I think we should unroll in cases like this, since I don't think Triton supports this either so that will be the only way to generate valid output code.