[Feature request] `T.ptr` cannot accept `nullptr` as input

[LyricZhao]

In some kernels, a tensor may be not used for some cases (e.g., with a specific template parameter), then users may pass None (nullptr) as the T.ptr. But ValueError: Unsupported tensor type: <class 'NoneType'> will be raised.

For example,

    @T.prim_func
    def main(
        a_ptr: T.ptr,
        b_ptr: T.ptr,
        c_ptr: T.ptr,
        bias_ptr: T.ptr,
        m: T.int32,
        n: T.int32,
        k: T.int32,
        with_bias: T.int32
    ):

with with_bias = 0, we expect bias_ptr as nullptr (but not forcing users to create a useless empty tensor).

[bucket-xv]

It is also often the case for T.Tensor like mask to have nullptr value. ValueError: Unsupported tensor type: <class 'NoneType'> will be raised. Below is an example:

 @T.prim_func
    def gate_kernel(
        mask: T.Tensor[(num_tokens, ), "bool"],
    ):

[LeiWang1999]

We can have a fix, but previously, we could achieve this goal using @tilelang.simplify_prim_func (as used in the TileLang-based library BitBLAS). For example:

with_bias = True

@tl.simplify_prim_func
def tl_matmul(
    A,
    B,
    C,
    Bias,
):
    ...
    if with_bias:
        f(Bias)

If with_bias is set to False, the program is automatically simplified to:

def tl_matmul(
    A,
    B,
    C,
):
    ...

In this simplified case, unused parameters are removed, so there is no need to pass the Bias tensor.

[LyricZhao]

@tilelang.simplify_prim_func seems not very intuitive & easy to use.

[LeiWang1999]

cc @xwhzz

[LeiWang1999]

now resolved, empty ptr we can check out the example test_tilelang_jit_nullptr.py