Python square brackets between function name and arguments: func[…](…)

Question:

I was learning how to accelerate python computations on GPU from this notebook, where one line confuses me:

mandel_kernel[griddim, blockdim](-2.0, 1.0, -1.0, 1.0, d_image, 20)

Here, mandel_kernel is a decorated (by cuda.jit) function, griddim and blockdim are tuples of length 2: griddim=(32,16), blockdim=(32,8).

Is this square brackets in between function name and argument list part of python syntax, or something specific to the cuda.jit decoration?

Asked By: Jason

||

Answers:

This is valid python syntax, i’ll try to break it down for you:

mandel_kernel is a dict, whose keys are 2-tuples (griddim, blockdim) and values are method (which is valid since methods are objects in python)

mandel_kernel[griddim, blockdim] therefore ‘returns’ (or evaluates to) a method

mandel_kernel[griddim, blockdim](-2.0, 1.0, -1.0, 1.0, d_image, 20) therefore calls that method with whatever arguments are inside the parenthesis.

This one line could be rewritten in three lines like so:

key = tuple(griddim, blockdim)
method = mandel_kernel[key]
method(-2.0, 1.0, -1.0, 1.0, d_image, 20)
Answered By: iCart

As iCart stated, @cuda.jit overloads the getitem method which changes the behavior of the brackets to instead of acting like a dict, the brackets are used to pass a tuple in a manner spiritually similar to the CUDA C/C++ runtime API syntax that looks like:

cudaKernel<<<numBlocks, threadsPerBlock>>>(parameter1, parameter2, …, parameterN);

It’s cute, but it’s anything but clear at first glance IMO. It’s doubly confusing when the tuple is computed by a lambda on the fly as it is for the triton library examples. I’m gathering they got this idea from numba.

Answered By: Scott Le Grand