AMGX in Julia
The AMGX.jl package provides an interface for using NVIDIA's AMGX library from the Julia language.
The package is installed using the Julia package manager:
using Pkg; Pkg.add("AMGX")
Currently, only prebuilt binaries are available for Linux, on other operating systems you need to have AMGX available locally.
Before using the package, reading through the official API reference docs for AMGX is recommended.
If you do not want to use the provided prebuilt binaries, set the environment variable
JULIA_AMGX_PATH to the path to the local AMGX library.
The library can now be initialized with:
using AMGX AMGX.initialize() AMGX.initialize_plugins()
Config can be created from a dictionary or a string:
config = AMGX.Config(Dict("monitor_residual" => 1, "max_iters" => 10, "store_res_history" => 1));
Resources object is created from an AMGX
Config. Currently, only simple resources are wrapped:
resources = AMGX.Resources(config)
The different modes in AMGX are available as:
AMGXVector is created from a resource object with a given mode.
v = AMGX.AMGXVector(resources, AMGX.dDDI)
Data can then be uploaded to the vectorusing
AMGX.upload!(v, [1.0, 2.0, 3.0])
Optionally, the "block dimension" can be given:
v_block = AMGX.AMGXVector(resources, AMGX.dDDI) AMGX.upload!(v_block, [1.0, 2.0, 3.0, 4.0]; block_dim=2)
Data can be downloaded from the vector (for example after solving a system) using
v_h = Vector(v)
It is also possible to download to a preallocated buffer using
v_h_buffer = zeros(3) copy!(v_h_buffer, v)
Note that data can be uploaded / downloaded from arrays already allocated on the GPU (
using CUDA v_cu = AMGX.AMGXVector(resources, AMGX.dDDI) cu = CuVector([1.0, 2.0, 3.0]) AMGX.upload!(v_cu, cu) cu_buffer = CUDA.zeros(Float64, 3) copy!(cu_buffer, v_cu)
A vector can be set to zero:
v_zero = AMGX.AMGXVector(resources, AMGX.dDDI) AMGX.set_zero!(v_zero, 5)
Matrices in AMGX are stored in CSR format (as opposed to CSC which is typically used in Julia).
AMGXMatrix is created from a
Resources and a
matrix = AMGX.AMGXMatrix(resources, AMGX.dDDI)
Data can be uploaded to it using the 3 CSR arrays:
AMGX.upload!(matrix, Cint[0, 1, 3], # row_ptrs Cint[1, 0, 1], # col_indices [1.0, 2.0, 3.0] # data )
These arrays can also be uploaded from
CuArrays already residing on the GPU.
CUDA.CUSPARSE.CuSparseMatrixCSR can be directly uploaded.
The non zero values can be replaced:
AMGX.replace_coefficients!(matrix, [3.0, 2.0, 1.0])
A solver is created from a
Mode, and a
solver = AMGX.Solver(resources, AMGX.dDDI, config)
A system can now be solved as:
x = AMGX.AMGXVector(resources, AMGX.dDDI) AMGX.set_zero!(x, 3) AMGX.setup!(solver, matrix) AMGX.solve!(x, solver, v)
The solution vector can now be downloaded:
After a solve, the status can be retrieved using
AMGX.get_status(solver). It is of type
AMGX.SolverStatus and can be either:
The total number of iterations can be retrieved with
The residual for a given iteration can be retrieved with
AMGX.get_iteration_residual(solver) AMGX.get_iteration_residual(solver, 0)
The API version is retrieved with
Some more version info can be printed using
For performance, it is recommended to pin host memory before uploading it to the GPU. Pinning and unpinning of memory is done using:
v = rand(5) AMGX.pin_memory(v) AMGX.unpin_memory(v)
By default, the AMGX library prints various things to
stdout. This can be overridden by registering a print callback, which is a Julia function accepting a
String and returning
str = "" store_to_str(amgx_printed::String) = (global str = amgx_printed; nothing) AMGX.register_print_callback(store_to_str) c_config = AMGX.Config("") print(str) print_stdout(amgx_printed::String) = print(stdout, amgx_printed) AMGX.register_print_callback(print_stdout)
Signal handlers can be installed and reset using:
Memory management and finalizing
You need to explicitly free memory of every AMGX object (
Solver) created using the julia call
Using Defer.jl can significantly
increase the convenience of this.
AMGX.jl contains a reference counting system so that it errors if you try to close things in the wrong order (e.g. closing the
Resources object before the
AMGXVector created from it is closed (destroyed))
When usage of the library is done, it should be finalized with:
The following functions from the C-API are not yet implemented:
AMGX_resources_create(only simple is currently wrapped)