Note that the 2d grid case is still incorrect. Consider: the grid size should be the number of compute units required, but the threadgroup size is a minimum number of compute units.
If I need to process a 50x50 image, I should be able to set a grid size of 50x50. Since the minimum threadgroup size is probably bigger (say 1024), this should turn into one dispatch of size (1, 1, 1). However with the current implementation, we will make a dispatch of (50, 50, 1), which essentially squares the amount of work - doing one thread group per unit of compute.
The correct implementation for 2d compute should take the mod of each grid dimension with the threadgroup size in that dimension. I did not fix this case as we do not have a use for 2d compute yet.
567f166b13