Annotating parameters with tl.pointer_type results in KeyError

[carmocca]

Annotating parameters with foo: tl.pointer_type results in the following stacktrace

my_model.py:90: in forward
    my_kernel[lambda meta: (triton.cdiv(numel, meta["BLOCK_SIZE"]),)](
../../../../.local/lib/python3.10/site-packages/triton/runtime/jit.py:345: in <lambda>
    return lambda *args, **kwargs: self.run(grid=grid, warmup=False, *args, **kwargs)
../../../../.local/lib/python3.10/site-packages/triton/runtime/autotuner.py:156: in run
    timings = {config: self._bench(*args, config=config, **kwargs) for config in pruned_configs}
../../../../.local/lib/python3.10/site-packages/triton/runtime/autotuner.py:156: in <dictcomp>
    timings = {config: self._bench(*args, config=config, **kwargs) for config in pruned_configs}
../../../../.local/lib/python3.10/site-packages/triton/runtime/autotuner.py:133: in _bench
    return do_bench(kernel_call, warmup=self.num_warmups, rep=self.num_reps, quantiles=(0.5, 0.2, 0.8))
../../../../.local/lib/python3.10/site-packages/triton/testing.py:103: in do_bench
    fn()
../../../../.local/lib/python3.10/site-packages/triton/runtime/autotuner.py:114: in kernel_call
    self.fn.run(
../../../../.local/lib/python3.10/site-packages/triton/runtime/jit.py:662: in run
    kernel = self.compile(
../../../../.local/lib/python3.10/site-packages/triton/compiler/compiler.py:276: in compile
    module = src.make_ir(options, codegen_fns, context)
../../../../.local/lib/python3.10/site-packages/triton/compiler/compiler.py:113: in make_ir
    return ast_to_ttir(self.fn, self, context=context, options=options, codegen_fns=codegen_fns)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
name = 'ier_type'
    def str_to_ty(name):
        if name[0] == "*":
            name = name[1:]
            if name[0] == "k":
                name = name[1:]
                ty = str_to_ty(name)
                return const_pointer_type(ty)
            ty = str_to_ty(name)
            return pointer_type(ty)
        tys = {
            "fp8e4nv": float8e4nv,
            "fp8e4b8": float8e4b8,
            "fp8e5": float8e5,
            "fp8e5b16": float8e5b16,
            "fp8e4b15": float8e4b15,
            "fp16": float16,
            "bf16": bfloat16,
            "fp32": float32,
            "fp64": float64,
            "i1": int1,
            "i8": int8,
            "i16": int16,
            "i32": int32,
            "i64": int64,
            "u1": int1,
            "u8": uint8,
            "u16": uint16,
            "u32": uint32,
            "u64": uint64,
            "B": int1,
        }
>       return tys[name]
E       KeyError: 'ier_type'
../../../../.local/lib/python3.10/site-packages/triton/language/__init__.py:[284]

The fix is simple: to remove these annotations.

But this used to work in previous versions.

Environment:

triton                    3.0.0
pytorch-triton            2.2.0+e28a256d71
torch                     2.4.1+cu118

Thanks!

[iclementine]

I do some debugging and find that the bug lies in the generated binder.

class JITFunction(KernelInterface[T]):
    def run(self, *args, grid, warmup, **kwargs):
        if self.binder is None:
            self.create_binder()
        ...
        bound_args, sig_and_spec, constexpr_vals, non_constexpr_vals, excess_kwargs = self.binder(*args, **kwargs)

The sig_and_spec returned from the binder is already wrong.

I made a jit function with the following signature

@triton.jit
def unary_pointwise_kernel(X: tl.pointer_type, O: tl.pointer_type, n: int, BLOCK_N: tl.constexpr):

the sig_and_spec is ('ier_type', 'ier_type', 'i32', 'D', 'D', 'D').

Finally I found that the bug is in

class KernelParam:
    ...
    @cached_property
    def annotation_type(self):
        annotation = self.annotation
        for ty1, ty2 in [("uint", 'u'), ("int", 'i')]:
            width = annotation[annotation.find(ty1) + len(ty1):]
            if width and ty1 in annotation:
                return f"{ty2}{width}"
        if annotation == "bool":
            return "u1"
        return ""

here annotatin is pointer_type (tl. is stripped beforehand), and int is in pointer_type. So width is er_type, the return value is ier_type.