[PERF] cuda.parallel: Cache intermediate results to improve performance of cudax.reduce_into

[shwina]

Description

This PR uses caching on the Python side to improve the performance of cuda.experimental.reduce_into. Specifically:

• cache _Reduce objects. The cache key used here is the dtype of the input arrays rather than the arrays themselves. I think this is safe to do, and longer term, I'd like to avoid passing the concrete arrays to the _Reduce constructor.

• cache the result of the utility function _type_to_info.

Before this PR:

In [4]: d_in = cuda.device_array(1, "int64")

In [5]: d_out = cuda.device_array(1, "int64")

In [6]: h_init = np.asarray([0], "int64")

In [7]: def op(x, y): return x + y

In [8]: %timeit cudax.reduce_into(d_in, d_out, op, h_init)
781 ms ± 42.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [9]: %timeit cudax._type_to_info(np.int32)
71.8 μs ± 111 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

After this PR:

In [8]: %timeit cudax.reduce_into(d_in, d_out, op, h_init)
The slowest run took 17.54 times longer than the fastest. This could mean that an intermediate result is being cached.
910 ns ± 1.36 μs per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [9]: %timeit cudax._type_to_info(np.int32)
64.8 ns ± 0.0356 ns per loop (mean ± std. dev. of 7 runs, 10,000,000 loops each)

Additional context

This came up in the initial investigations for #2958.

Checklist

• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[miscco]

Welcome to the team 🎉

[miscco]

/ok to test

[github-actions]

🟩 CI finished in 15m 30s: Pass: 100%/1 | Total: 15m 30s | Avg: 15m 30s | Max: 15m 30s

• 🟩 python: Pass: 100%/1 | Total: 15m 30s | Avg: 15m 30s | Max: 15m 30s

  🟩 cpu
    🟩 amd64              Pass: 100%/1   | Total: 15m 30s | Avg: 15m 30s | Max: 15m 30s
  🟩 ctk
    🟩 12.6               Pass: 100%/1   | Total: 15m 30s | Avg: 15m 30s | Max: 15m 30s
  🟩 cudacxx
    🟩 nvcc12.6           Pass: 100%/1   | Total: 15m 30s | Avg: 15m 30s | Max: 15m 30s
  🟩 cudacxx_family
    🟩 nvcc               Pass: 100%/1   | Total: 15m 30s | Avg: 15m 30s | Max: 15m 30s
  🟩 cxx
    🟩 GCC13              Pass: 100%/1   | Total: 15m 30s | Avg: 15m 30s | Max: 15m 30s
  🟩 cxx_family
    🟩 GCC                Pass: 100%/1   | Total: 15m 30s | Avg: 15m 30s | Max: 15m 30s
  🟩 gpu
    🟩 v100               Pass: 100%/1   | Total: 15m 30s | Avg: 15m 30s | Max: 15m 30s
  🟩 jobs
    🟩 Test               Pass: 100%/1   | Total: 15m 30s | Avg: 15m 30s | Max: 15m 30s
  

👃 Inspect Changes

Modifications in project?

	Project
	CCCL Infrastructure
	libcu++
	CUB
	Thrust
	CUDA Experimental
+/-	python
	CCCL C Parallel Library
	Catch2Helper

Modifications in project or dependencies?

	Project
	CCCL Infrastructure
	libcu++
	CUB
	Thrust
	CUDA Experimental
+/-	python
	CCCL C Parallel Library
	Catch2Helper

🏃‍ Runner counts (total jobs: 1)

#	Runner
1	linux-amd64-gpu-v100-latest-1

[shwina]

Looks like the build docs CI job is using Python=3.7 which is missing the functools.cache function. Looking into it.

[github-actions]

🟩 CI finished in 14m 53s: Pass: 100%/1 | Total: 14m 53s | Avg: 14m 53s | Max: 14m 53s

• 🟩 python: Pass: 100%/1 | Total: 14m 53s | Avg: 14m 53s | Max: 14m 53s

  🟩 cpu
    🟩 amd64              Pass: 100%/1   | Total: 14m 53s | Avg: 14m 53s | Max: 14m 53s
  🟩 ctk
    🟩 12.6               Pass: 100%/1   | Total: 14m 53s | Avg: 14m 53s | Max: 14m 53s
  🟩 cudacxx
    🟩 nvcc12.6           Pass: 100%/1   | Total: 14m 53s | Avg: 14m 53s | Max: 14m 53s
  🟩 cudacxx_family
    🟩 nvcc               Pass: 100%/1   | Total: 14m 53s | Avg: 14m 53s | Max: 14m 53s
  🟩 cxx
    🟩 GCC13              Pass: 100%/1   | Total: 14m 53s | Avg: 14m 53s | Max: 14m 53s
  🟩 cxx_family
    🟩 GCC                Pass: 100%/1   | Total: 14m 53s | Avg: 14m 53s | Max: 14m 53s
  🟩 gpu
    🟩 v100               Pass: 100%/1   | Total: 14m 53s | Avg: 14m 53s | Max: 14m 53s
  🟩 jobs
    🟩 Test               Pass: 100%/1   | Total: 14m 53s | Avg: 14m 53s | Max: 14m 53s
  

👃 Inspect Changes

Modifications in project?

	Project
	CCCL Infrastructure
	libcu++
	CUB
	Thrust
	CUDA Experimental
+/-	python
	CCCL C Parallel Library
	Catch2Helper

Modifications in project or dependencies?

	Project
	CCCL Infrastructure
	libcu++
	CUB
	Thrust
	CUDA Experimental
+/-	python
	CCCL C Parallel Library
	Catch2Helper

🏃‍ Runner counts (total jobs: 1)

#	Runner
1	linux-amd64-gpu-v100-latest-1

[shwina]

Looks like the build docs CI job is using Python=3.7 which is missing the functools.cache function. Looking into it.

Punting on this and just using @functools.lru_cache instead which is supported by Python 3.7.

[rwgk]

That's what I was wondering about, but I didn't get to drilling down.

It might be useful to work together on getting this TODO done.

Already, this PR will have complicated merge conflicts with my #2788, i.e. it might be best to team up working on both.

[shwina]

It might be useful to work together on getting this TODO done.

Definitely! I opened #3008 to track this. Let's tackle it as a follow up to this PR.