AMGX in Julia
Build Status |
---|
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()
An AMGX Config
can be created from a dictionary or a string:
config = AMGX.Config(Dict("monitor_residual" => 1, "max_iters" => 10, "store_res_history" => 1));
An AMGX 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:
AMGX.hDDI
AMGX.hDFI
AMGX.hFFI
AMGX.dDDI
AMGX.dDFI
AMGX.dFFI
An 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 upload!
:
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 Vector
:
v_h = Vector(v)
It is also possible to download to a preallocated buffer using copy!
:
v_h_buffer = zeros(3)
copy!(v_h_buffer, v)
Note that data can be uploaded / downloaded from arrays already allocated on the GPU (CuArray
):
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).
An AMGXMatrix
is created from a Resources
and a mode
:
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.
Alternatively, a 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 Resources
, a Mode
, and a Config
:
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:
Vector(x)
After a solve, the status can be retrieved using AMGX.get_status(solver)
. It is of type AMGX.SolverStatus
and can be either:
AMGX.SUCCESS
AMGX.FAILED
AMGX.DIVERGED
The total number of iterations can be retrieved with AMGX.get_iterations_number(solver)
.
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 AMGX.api_version()
.
Some more version info can be printed using AMGX.versioninfo()
.
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 nothing
:
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:
AMGX.install_signal_handler()
AMGX.reset_signal_handler()
You need to explicitly free memory of every AMGX object (Config
, Resources
, AMGXVector
, AMGXMatrix
, Solver
) created using the julia call
close
.
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:
AMGX.finalize_plugins()
AMGX.finalize()
The following functions from the C-API are not yet implemented:
AMGX_read_system
AMGX_read_system_distributed
AMGX_write_system
AMGX_write_system_distributed
AMGX_config_create_from_file
AMGX_config_get_default_number_of_rings
AMGX_resources_create
(only simple is currently wrapped)AMGX_matrix_upload_all_global
AMGX_matrix_comm_from_maps
AMGX_matrix_comm_from_maps_one_ring
AMGX_vector_bind