Skip to content

Commit

Permalink
Add support for multiple GPUs in echos.jl
Browse files Browse the repository at this point in the history
  • Loading branch information
@stijnh
stijnh committed Dec 16, 2024
1 parent 3b67aa1 commit b3dd94d
Show file tree
Hide file tree
Showing 18 changed files with 227 additions and 147 deletions.
8 changes: 8 additions & 0 deletions CompasToolkit.jl/src/CompasToolkit.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,14 @@ function get_context()::Context
end
end

"""
Waits until all asynchronous operations of Compas have finished.
"""
function synchronize()
context = get_context()
@ccall LIBRARY.compas_synchronize(pointer(context)::Ptr{Cvoid})::Cvoid
end

"""
Object representing an array of size `N` and type `T`.
"""
Expand Down
25 changes: 25 additions & 0 deletions include/compas/core/context.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,27 @@ struct CompasContext {
return m_runtime.allocate(content.data(), content.size());
}

template<typename F, typename... Args>
void
parallel_device(kmm::NDRange index_space, kmm::NDSize chunk_size, F fun, Args... args) const {
m_runtime
.parallel_submit(index_space, kmm::TaskPartitioner(chunk_size), kmm::GPU(fun), args...);
}

template<typename F, typename... Args>
void parallel_kernel(
kmm::NDRange index_space,
kmm::NDSize chunk_size,
dim3 block_dim,
F kernel,
Args... args) const {
m_runtime.parallel_submit(
index_space,
kmm::TaskPartitioner(chunk_size),
kmm::GPUKernel(kernel, block_dim),
args...);
}

template<typename F, typename... Args>
void submit_device(kmm::NDRange index_space, F fun, Args... args) const {
m_runtime.submit(index_space, m_device, kmm::GPU(fun), args...);
Expand All @@ -58,6 +79,10 @@ struct CompasContext {
m_runtime.synchronize();
}

const kmm::Runtime& runtime() const {
return m_runtime;
}

private:
kmm::Runtime m_runtime;
kmm::DeviceId m_device;
Expand Down
37 changes: 17 additions & 20 deletions include/compas/core/view.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,31 +5,28 @@ namespace compas {

using index_t = int;

template<typename T, size_t N = 1>
using view = kmm::view<T, N>;

template<typename T, size_t N = 1>
using view_mut = kmm::view_mut<T, N>;
using kmm::strided_subview;
using kmm::strided_subview_mut;
using kmm::strided_view;
using kmm::strided_view_mut;
using kmm::subview;
using kmm::subview_mut;
using kmm::view;
using kmm::view_mut;

using kmm::gpu_strided_subview;
using kmm::gpu_strided_subview_mut;
using kmm::gpu_strided_view;
using kmm::gpu_strided_view_mut;
using kmm::gpu_subview;
using kmm::gpu_subview_mut;
using kmm::gpu_view;
using kmm::gpu_view_mut;

template<typename T, size_t N = 1>
using host_view = kmm::view<T, N>;

template<typename T, size_t N = 1>
using host_view_mut = kmm::view_mut<T, N>;

template<typename T, size_t N = 1>
using gpu_view = kmm::gpu_view<T, N>;

template<typename T, size_t N = 1>
using gpu_view_mut = kmm::gpu_view_mut<T, N>;

template<typename T, size_t N = 1>
using gpu_subview_mut = kmm::gpu_subview_mut<T, N>;

template<typename T, size_t N = 1>
using gpu_strided_view = kmm::gpu_strided_view<T, N>;

template<typename T, size_t N = 1>
using gpu_strided_view_mut = kmm::gpu_strided_view_mut<T, N>;

} // namespace compas
31 changes: 14 additions & 17 deletions include/compas/parameters/tissue.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,20 +13,17 @@ namespace compas {
* column is a voxel.
*/
struct TissueParameters: public Object {
Array<float, 2> parameters; // Size: [TissueParameterField::NUM_FIELDS, nvoxels]
Array<float, 2> data; // Size: [TissueParameterField::NUM_FIELDS, nvoxels]
int nvoxels;
int chunk_size;
bool has_z = true;
bool has_b0 = true;
bool has_b1 = true;

TissueParameters(
Array<float, 2> parameters,
int nvoxels,
bool has_z,
bool has_b0,
bool has_b1) :
parameters(parameters),
nvoxels(nvoxels),
TissueParameters(Array<float, 2> parameters, bool has_z, bool has_b0, bool has_b1) :
data(parameters),
nvoxels(kmm::checked_cast<int>(parameters.size(1))),
chunk_size(kmm::checked_cast<int>(parameters.chunk_size(1))),
has_z(has_z),
has_b0(has_b0),
has_b1(has_b1) {}
Expand All @@ -35,6 +32,7 @@ struct TissueParameters: public Object {
TissueParameters make_tissue_parameters(
const CompasContext& ctx,
int num_voxels,
int chunk_size,
view<float> T1,
view<float> T2,
view<float> B1,
Expand All @@ -43,7 +41,7 @@ TissueParameters make_tissue_parameters(
view<float> rho_y,
view<float> x,
view<float> y,
view<float> z);
view<float> z = {});

} // namespace compas

Expand All @@ -53,13 +51,12 @@ struct Argument<compas::TissueParameters> {
using type = compas::TissueParametersView;

static Argument pack(TaskBuilder& builder, compas::TissueParameters p) {
return {
{.parameters = {}, //
.nvoxels = p.nvoxels,
.has_z = p.has_z,
.has_b0 = p.has_b0,
.has_b1 = p.has_b1},
pack_argument(builder, p.parameters)};
compas::TissueParametersView view;
view.has_z = p.has_z;
view.has_b0 = p.has_b0;
view.has_b1 = p.has_b1;

return {view, pack_argument(builder, p.data)};
}

template<ExecutionSpace Space>
Expand Down
11 changes: 6 additions & 5 deletions include/compas/parameters/tissue_view.cuh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,11 +39,12 @@ struct TissueVoxel {
* Device-side representation of `TissueParameters`
*/
struct TissueParametersView {
gpu_view<float, 2> parameters;
int nvoxels;
bool has_z;
bool has_b0;
bool has_b1;
TissueParametersView(gpu_subview<float, 2> parameters = {}) : parameters(parameters) {}

gpu_subview<float, 2> parameters;
bool has_z = true;
bool has_b0 = true;
bool has_b1 = true;

/**
* Returns the parameters for the voxel at location `i`.
Expand Down
8 changes: 8 additions & 0 deletions julia-bindings/src/bindings.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,10 @@ extern "C" void compas_destroy_context(const compas::CompasContext* ctx) {
return catch_exceptions([&] { delete ctx; });
}

extern "C" void compas_synchronize(const compas::CompasContext* ctx) {
return catch_exceptions([&] { ctx->synchronize(); });
}

extern "C" const kmm::ArrayBase* compas_make_array_float(
const compas::CompasContext* context,
const float* data_ptr,
Expand Down Expand Up @@ -130,10 +134,14 @@ extern "C" const compas::TissueParameters* compas_make_tissue_parameters(
const float* x,
const float* y,
const float* z) {
int num_devices = int(context->runtime().info().num_devices());
int chunk_size = kmm::round_up_to_multiple(kmm::div_ceil(nvoxels, num_devices), 32);

return catch_exceptions([&] {
auto params = compas::make_tissue_parameters(
*context,
nvoxels,
chunk_size,
make_view(T1, nvoxels),
make_view(T2, nvoxels),
make_view(B1, nvoxels),
Expand Down
4 changes: 2 additions & 2 deletions src/jacobian/hermitian.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -74,8 +74,8 @@ Array<cfloat, 2> compute_jacobian_hermitian(
echos, \
delta_echos_T1, \
delta_echos_T2, \
parameters.parameters.size(1), \
parameters.parameters, \
parameters.data.size(1), \
parameters.data, \
coil_sensitivities, \
vector, \
E, \
Expand Down
4 changes: 2 additions & 2 deletions src/jacobian/product.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -68,8 +68,8 @@ Array<cfloat, 3> compute_jacobian(
echos, \
delta_echos_T1, \
delta_echos_T2, \
parameters.parameters.size(1), \
parameters.parameters, \
parameters.data.size(1), \
parameters.data, \
coil_sensitivities, \
E, \
dEdT2, \
Expand Down
66 changes: 35 additions & 31 deletions src/parameters/tissue.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ namespace compas {
TissueParameters make_tissue_parameters(
const CompasContext& ctx,
int num_voxels,
int voxels_per_chunk,
host_view<float> T1,
host_view<float> T2,
host_view<float> B1,
Expand All @@ -15,70 +16,73 @@ TissueParameters make_tissue_parameters(
host_view<float> x,
host_view<float> y,
host_view<float> z) {
auto stride = round_up_to_multiple_of(num_voxels, 32);
auto params = kmm::Array<float, 2> {{TissueParameterField::NUM_FIELDS, stride}};
using namespace kmm::placeholders;
auto params = kmm::Array<float, 2> {{TissueParameterField::NUM_FIELDS, num_voxels}};

bool has_z = !z.is_empty();
bool has_b0 = !B0.is_empty();
bool has_b1 = !B1.is_empty();

ctx.submit_device(
ctx.parallel_device(
num_voxels,
[&](kmm::DeviceContext& device, kmm::NDRange, gpu_view_mut<float, 2> params) {
device.fill(params, 0.0F);
voxels_per_chunk,
[&](kmm::DeviceContext& device, kmm::NDRange range, gpu_subview_mut<float, 2> params) {
KMM_ASSERT(params.is_contiguous());
device.fill(params.data(), params.size(), 0.0F);
auto offset = range.begin();
auto length = range.size();

device.copy(
T1.data(),
params.drop_axis<0>(TissueParameterField::T1).data(),
num_voxels);
T1.data_at(offset),
params.data_at(TissueParameterField::T1, offset),
length);
device.copy(
T2.data(),
params.drop_axis<0>(TissueParameterField::T2).data(),
num_voxels);
T2.data_at(offset),
params.data_at(TissueParameterField::T2, offset),
length);

device.copy(
rho_x.data(),
params.drop_axis<0>(TissueParameterField::RHO_X).data(),
num_voxels);
rho_x.data_at(offset),
params.data_at(TissueParameterField::RHO_X, offset),
length);
device.copy(
rho_y.data(),
params.drop_axis<0>(TissueParameterField::RHO_Y).data(),
num_voxels);
rho_y.data_at(offset),
params.data_at(TissueParameterField::RHO_Y, offset),
length);

device.copy(x.data(), params.drop_axis<0>(TissueParameterField::X).data(), num_voxels);
device.copy(y.data(), params.drop_axis<0>(TissueParameterField::Y).data(), num_voxels);
device.copy(x.data_at(offset), params.data_at(TissueParameterField::X, offset), length);
device.copy(y.data_at(offset), params.data_at(TissueParameterField::Y, offset), length);

if (has_z) {
device.copy(
z.data(),
params.drop_axis<0>(TissueParameterField::Z).data(),
num_voxels);
z.data_at(offset),
params.data_at(TissueParameterField::Z, offset),
length);
}

if (has_b0) {
device.copy(
B0.data(),
params.drop_axis<0>(TissueParameterField::B0).data(),
num_voxels);
B0.data_at(offset),
params.data_at(TissueParameterField::B0, offset),
length);
}

if (has_b1) {
device.copy(
B1.data(),
params.drop_axis<0>(TissueParameterField::B1).data(),
num_voxels);
B1.data_at(offset),
params.data_at(TissueParameterField::B1, offset),
length);
} else {
// The default value for B1 is 1
device.fill(params.drop_axis<0>(TissueParameterField::B1), 1.0f);
device.fill(params.data_at(TissueParameterField::B1, offset), length, 1.0F);
}
},
write(params));
write(params, access(_, _x)));

params.synchronize();

return {
params,
num_voxels,
has_z,
has_b0,
has_b1,
Expand Down
Loading

0 comments on commit b3dd94d

Please sign in to comment.