The thread extents are part of the params of the PrimFunc for the device code, kernel launch parameters are set here c(https://github.com/apache/tvm/blob/main/src/tir/transforms/split_host_device.cc#L294-L296). In LaunchParamConfig, base and arg_index_map is used to map the index of params of PrimFunc to the thread extent which is an array of six integers representing (blockIdx.x/y/z, threadIdx.x/y/z).
work_dim is the dimension of the thread extent, which is the highest dimension among x/y/z specified in the thread extent
2 Likes