[Bug] Buggy with empty/dummy kernels

[LyricZhao]

For some cases, we want to use dummy kernels to pad the CUDA stream for specific purpose (e.g. to delay the later kernel launches). But for the following cases, TileLang fails to compile an empty kernel:

pid length matching

import tilelang
from tilelang import language as T


@tilelang.jit()
def get_buggy_kernel():
    @T.prim_func
    def buggy_kernel():
        with T.Kernel(1, threads=32) as (pid, ):
            pass

    return buggy_kernel


if __name__ == '__main__':
    kernel = get_buggy_kernel()
    print(kernel.get_kernel_source())
    kernel()

error: cannot unpack non-iterable Var object, which is not friendly and I struggle to find (pid, ) has to be changed into pid.

Failed to compile

import tilelang
from tilelang import language as T


@tilelang.jit()
def get_buggy_kernel():
    @T.prim_func
    def buggy_kernel():
        with T.Kernel(1, threads=32) as pid:
            pass

    return buggy_kernel


if __name__ == '__main__':
    kernel = get_buggy_kernel()
    print(kernel.get_kernel_source())
    kernel()

TypeError: Downcast from ir.IntImm to tir.Call failed.

[LyricZhao]

Another case:

import torch
import tilelang
from tilelang import language as T


@tilelang.jit()
def get_buggy_kernel():
    num_tokens = T.symbolic('num_tokens')

    @T.prim_func
    def buggy_kernel(x: T.Tensor[(num_tokens, ), 'float'],):
        with T.Kernel(num_tokens, threads=32) as pid:
            y = x[pid]

    return buggy_kernel


if __name__ == '__main__':
    kernel = get_buggy_kernel()
    print(kernel.get_kernel_source())

    x = torch.randn((128, ), dtype=torch.float, device='cuda')
    kernel(x)

Debugging this just want to see whether x[pid] will IMA, while TileLang optimizes the whole block into empty and leads to compilation failure as the previous cases.

[SiriusNEO]

TypeError: Downcast from ir.IntImm to tir.Call failed. this is because TVM usually uses T.evaluate(0) as dummy op and a customized pass in Tilelang doesn't handle this case (Directly cast args of T.evaluate to Call): https://github.com/tile-ai/tilelang/blob/main/src/transform/lower_tile_op.cc#L88

But after fix it, I find that SplitHostDevice fails to split dummy kernel to device (There is no function in the device mod). Then it is checked by this assertion: https://github.com/tile-ai/tilelang/blob/main/tilelang/jit/adapter/wrapper.py#L457. A refactor in the SplitHostDevice strategy may be needed to run simple kernels in GPU. cc @LeiWang1999

[LeiWang1999]

TypeError: Downcast from ir.IntImm to tir.Call failed. this is because TVM usually uses T.evaluate(0) as dummy op and a customized pass in Tilelang doesn't handle this case (Directly cast args of T.evaluate to Call): https://github.com/tile-ai/tilelang/blob/main/src/transform/lower_tile_op.cc#L88

But after fix it, I find that SplitHostDevice fails to split dummy kernel to device (There is no function in the device mod). Then it is checked by this assertion: https://github.com/tile-ai/tilelang/blob/main/tilelang/jit/adapter/wrapper.py#L457. A refactor in the SplitHostDevice strategy may be needed to run simple kernels in GPU. cc @LeiWang1999

It's likely that pr #962 addressed the first part of this comments, but the second one about SplitHostDevice is a bit complicated, @Hzfengsy do you have any suggestions about this pass?

[LeiWang1999]

Resolved, as PR #962 has been merged. Now, [this test file](https://github.com/tile-ai/tilelang/blob/main/testing/python/issue/test_tilelang_issue_830.py) includes all the tests we've encountered for this issue.

Specifically:

if not hasattr(Var, "__iter__"):

    def _var_iter(self):
        yield self

    Var.__iter__ = _var_iter  # type: ignore[attr-defined]

if not hasattr(Var, "__len__"):
    Var.__len__ = lambda self: 1  # type: ignore[attr-defined]

Now, both of these code styles are valid and will return the correct pid value.

@tilelang.jit
def _empty_kernel_with_binding_variants(use_tuple_binding: bool = False):

    @T.prim_func
    def kernel_with_tuple_kernel_binding():
        with T.Kernel(1, threads=32) as (pid,):
            print(pid)
            pass

    @T.prim_func
    def kernel_with_scalar_kernel_binding():
        with T.Kernel(1, threads=32) as pid:
            print(pid)
            pass

    return kernel_with_tuple_kernel_binding if use_tuple_binding else kernel_with_scalar_kernel_binding

[LeiWang1999]

closed and thanks @SiriusNEO !