Add GPU Functionality to the dev Branch

[braxtoncuneo]

Summary of changes

This PR adds GPU execution to the dev branch. This was accomplished by re-implementing some functionality in a manner compatible with execution through numba-hip and numba-cuda. This includes changes to how gpu state objects are managed, the introduction of helper intrinsics/functions to handle returning arrays, and replacement of builtin functions that are not supported on one or more gpu plaltforms (e.g. powers for complex numbers, numpy polynomial solvers).

Types of changes

New feature: Helper functions/intrinsics for returning arrays on GPU

As has been mentioned in previous issues/PRs, numba-hip and numba-cuda do not permit returning arrays from functions. In brief, this is because memory management on GPU is hard, and the simplest way to create an array as a local variable is to allocate it on the stack. Once the array that created a local array returns, that array no longer "exists" in a way that it can be safely used.

This is a problem, since the helper functions provided by numba_objects_generator.py return references to and slices of arrays in the mcdc global data structure. We'd like to keep those helper functions, considering how much convenience and quality of life they provide to developers. Also, since these global data structures have a lifetime which encompasses the lifetime of all transport logic (i.e. the structures are in memory and safe to use whenever any transport logic is run) it should be safe to ignore this array-returning restriction for these helper functions.

For this specific use case, when we know for certain that an array is safe to use after a function returns, this PR introduces the array_return decorator and the array_result function, which enables the return of array references/slices like so:

# The decorator must be supplied the type of element stored
# by the arrays that are returned by the decorated function.
# 
# An arbitrary number of parameters may be accepted by an
# array-returning function, as long as that number is fixed
# for a given array returning function.
@array_return(nb.types.float64)
def example_function(param_1,param_2,param_3):
    # ...array slicing logic...

    # Returned arrays must always be wrapped in the
    # array_result function.
    return array_result(example_result_array_here)

This works by smuggling a pointer to the array's data and the size of the array out as a tuple, then reconstructing the array based off of the returned tuple in the calling function.

In normal python execution, the array_result function simply returns its input. However, there exists an overload of array_result which will handle the conversion to tuple in compiled contexts.

The array_return decorator registers an intrinsic overload of the decorated function. This intrinsic calls the decorated function and creates a new array based on the returned tuple. Because intrinsics are inlined directly into their calling function, they aren't their own function with a separate stack frame and so are not rejected by the array-returning guards of numba-hip or numba-cuda.

New feature: GPU-compatible quartic solver and complex power calculation

• Good news: The torus has been added as a new geometric primitive to MC/DC!
• Bad news: Solving torus/line intersections analytically requires solving quartic functions.
• Good news: numpy.roots solves polynomials and is implemented by numba in compiled contexts!
• Bad news: numba-hip and numba-cuda don't implement numpy.roots - likely because the general case does not have analytic solutions and requires the use of arbitrarily large matrixes (determined by the degree of the polynomial).
• Good news: Quartic functions do have an analytic solution!
• Bad news: Quartic functions require the calculation of real-valued powers for complex numbers, and numba-hip doesn't implement the intrinsic for powers of complex numbers.
• Good news: Integer powers of complex numbers can be calculated through iterative multiplication, complex square roots can be solved with square roots of real numbers, and fractional powers of complex numbers can be calculated using de Moivre's Formula!

@njit()
def modulus(x):
    return math.sqrt(x.real**2 + x.imag**2)


@njit()
def sqrt(x):
    return math.sqrt(x.real) * (x + x.real) / modulus(x + x.real)


@njit()
def power(x, n):
    result = 1
    for i in range(n):
        result = result * x
    return result


@njit()
def nth_root(x, n, index):
    # First, convert to polar form
    r = modulus(x)
    a = math.atan2(x.imag, x.real)

    # Apply de Moivre's Formula
    root_modulus = math.pow(r, 1.0 / n)
    root_argument = (a + 2 * math.pi * index) / n

    # ...then convert back to rectangular form
    real = root_modulus * math.cos(root_argument)
    imag = root_modulus * math.sin(root_argument)
    return complex(real, imag)


@njit()
def principal_nth_root(x, n):
    return nth_root(x, n, 0)

New feature: GPU execution

New feature: Added logic to test/regression/conftest.py to support regression/units tests with GPU

Bug fix: updated references to reflect the new names and module locations of Harmonize API, as imported by MC/DC

Bug fix: updated logic switched by gpu-related settings to check values in settings based upon the constants in mcdc.constant

Bug fix: removed redundant call to clear_flags in code_factory/gpu/transport/simulation.py

Bug fix: ensured that the original value of atomically-modified fields is taken from the return of the atomic operation, rather than from a separate read. This is important, because only the return value of the atomic operation itself can be trusted with providing the original value immediately before the operation.

Bug fix: updated add_bank_size and its uses so that they follow the interface of atomic operations, returning the trusted value. Using a separate get_bank_size after the fact has the same problem as a separate read, and so the original value must be provided by the logic of add_bank_size itself.

Bug fix: modified the reading of particles from banks so that they use the trusted values returned by add_bank_size, which requires the movement to occur after the addition itself has occurred.

Bug fix: added early return to progress-printing code in mcdc/transport/simulation.py to avoid divide-by-zero when mpi_work_size==0

Bug fix: updated the cpu-side atomic_add to match the correct atomic operation interface

Refactor: Added gpu platform detection to config.py, providing global flags to indicate whether CUDA/ROCM is available

Bug fix: Moved the call to setup_gpu_program in main so that it occurs after types have been defined. Without this ordering, the function would attempt to compile functions without proper type information.

Bug fix: Updated test/unit/test_numba_layers_generator.py to reflect the new array-returning functionality provided for mcdc_get/mcdc_set.

Developer Checklist

• I have read the contributing guide.
• My code follows the code style of this project.
• All new and existing tests pass

Associated Issues and PRs

• related to #

Associated Developers

[braxtoncuneo]

It should be noted that the slab_isobeam_td_census example was removed from the regression tests, which failed close-to-identically on both python/numba-mode and both with cpu or gpu as the target.

After consultation with @ilhamv , this seemed like the best option to get gpu functionality merged in a timely manner.

[braxtoncuneo]

These latest changes remove dead code, and so should not affect regression/unit tests.

[ilhamv]

Thanks, @braxtoncuneo! I'll work on the review ASAP. And thanks for highlighting some of the changes in separate conversations.

[ilhamv]

Dedicated pages in docs would be needed to reflect the updated GPU support. We will do that in a follow-up PR.

[ilhamv]

We would need to keep track of overheads that appear only in pure Python mode, as it would be important in interpreting runtime results when we compare Python VS Numba-CPU modes. We can do this when we work on the docs in the follow-up PR.

[ilhamv]

@braxtoncuneo

setup_gpu_program is called twice, first in generate_numba_layers() then in prepare(). Is that on purpose?

[ilhamv]

@braxtoncuneo
I think I fixed the bug in slab_isobeam_td_census. The issue was in using the returned value of add_bank_size when we promote a future particle to a census particle.

[ilhamv]

@braxtoncuneo
If now mcdc_get supports Numba-compatible array returns, should we "release" all the suppressed instances of "Above is equivalent to:"?

[ilhamv]

@braxtoncuneo
Note that I organize the root solver, renaming geometry/root_solve.py into geometry/surface/torus_root_solver.py, as it is currently used exclusively for torus.

I can see that there is potential for generalizing it to be a utility module. However, generalizing it would need to make the functions tangent-safe, like those currently used in the quadratic surfaces, not to mention the sacrificed performance from moving away from the specialized solvers. Nevertheless, due to the potential improved maintainability, we keep the generalization option open in the future!

[braxtoncuneo]

@braxtoncuneo If now mcdc_get supports Numba-compatible array returns, should we "release" all the suppressed instances of "Above is equivalent to:"?

You're right. I'll get those switched back to the mcdc_get equivalents.

[braxtoncuneo]

@braxtoncuneo

setup_gpu_program is called twice, first in generate_numba_layers() then in prepare(). Is that on purpose?

No. In fact, it looks like some of the changes between the CEMeNT dev and the MCDC project dev already fixed the ordering problem for setup_gpu_program and defining types. This latest batch of commits removes the additional call.

@braxtoncuneo I think I fixed the bug in slab_isobeam_td_census. The issue was in using the returned value of add_bank_size when we promote a future particle to a census particle.

I have verified that it passes for gpu on tuolumne. 👍

@braxtoncuneo If now mcdc_get supports Numba-compatible array returns, should we "release" all the suppressed instances of "Above is equivalent to:"?

The latest batch of commits revises these accesses back to the equivalent mcdc_get uses.