[Do not merge] Spike upgrading comps algorithm with taichi

[jeancochrane]

⚠️ I don't plan to merge this PR, but am requesting review for the sake of knowledge transfer. ⚠️

This PR tests out the use of taichi for the comps algorithm rather than numba.

Connects #229.

Findings

See Benchmarking below for detailed stats comparing this approach to our existing code on different architectures. Some high-level takeaways:

• CUDA doesn't seem to make much of a difference, and is counterproductive if anything. This makes me wonder whether the algorithm needs to be redesigned to make better use of the GPU (note that numba has an entirely separate interface for CUDA programming), but I'm considering that question out of scope for now.
• There are big performance gains to be had by simply bumping the instance type with the existing numba code. If the numbers below hold, we could speed up the comps code by 2x if we switched to c5.24xlarge instances, which are about twice as expensive as the m4.10xlarge instances we use now (meaning we should expect to break even on the change).
• At small scales (20k observations/10k comparisons), taichi appears to outperform numba, but this improvement disappears if we scale up the size of the data. At a large scale (100k observations/50k comparisons), they perform about the same.
• As evidenced by the code in this PR, the taichi interface is harder to work with than numba. Taichi is more strict about types, and doesn't support some basic operations that Python does (most notably, getting the shape of an array and returning an array from a function) making the code more confusing and un-Pythonic.

Benchmarking

20k observations, 10k comparisons

framework	instance type	arch	time	logs
taichi	g5.12xlarge	x86	2.36s	link
taichi	g5.12xlarge	CUDA	4.33s	link
taichi	m4.10xlarge	x86	4.44s	link
numba	g5.12xlarge	x86	6.07s	link
numba	m4.10xlarge	x86	10.52s	link

100k observations, 50k comparisons

framework	instance type	arch	time	logs
numba	g5.12xlarge	x86	31.87s	link
taichi	c5.24xlarge	x86	31.93s	link
taichi	m4.10xlarge	x86	34.09s	link
numba	c5.24xlarge	x86	37.31s	link
taichi	g5.12xlarge	x86	37.75s	link
taichi	g5.12xlarge	CUDA	43.58s	link
numba	m4.10xlarge	x86	64.19s	link

[jeancochrane]

Installing CUDA appears to be a pain, so for the purposes of benchmarking this PR I just switched to a base image that comes with CUDA installed and refactored the Dockerfile to only install dependencies necessary for running python/comps.py.

[jeancochrane]

shape and len attributes are not available inside taichi kernels, so the recommended approach for these types of array metadata is to compute them outside the kernel and then pass them into the kernel as arguments.

[jeancochrane]

Taichi cannot automatically convert numpy arrays to the data structures that it works with (called fields) so we have to do that manually prior to passing them into the kernel.

[jeancochrane]

Taichi also cannot return fields from kernels, so we need to define the output data structure ahead of time and then mutate it in the kernel function.

[jeancochrane]

for/else constructs are not supported in taichi, so we have to use a flag variable instead.

[jeancochrane]

Arrays cannot be defined inside a taichi kernel, and must instead be defined outside its scope and then passed in as an argument; hence, we needed to refactor this function to remove the dynamic allocation of arrays inside the context of the function.

[wagnerlmichael]

Nice catch on the instance bump! Thanks for sharing your findings. I wonder if mentioning the CUDA interface to the ex-intern who might take a look at this would be helpful, although I suspect their class is teaching CUDA interface programming.

I would also be happy to try the CUDA interface, I worked with it a bit in graduate school. I think trying CuPy could also be worthwhile.

[dfsnow]

Super cool work. Sad that it didn't speed things up. Sorry for sending you down the rabbit hole. Let's leave this up as I'd like to test it with a local GPU as well just to make sure we're using full GPU capacity.