Skip to content

fix(gpu): support LibDDLA 0.0.6 public API #19

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
PeizeLin merged 1 commit into from
Nov 10, 2026
Merged

fix(gpu): support LibDDLA 0.0.6 public API #19

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
87 changes: 42 additions & 45 deletions include/RI/global/gpu/GPU_Backend.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,9 +15,7 @@
#if (defined(DDLA_USE_CUDA) + defined(DDLA_USE_HIP)) != 1
#error "__DDLA_RI requires exactly one of DDLA_USE_CUDA or DDLA_USE_HIP"
#endif
#include <ddla/ddla_connector.h>
#include <ddla/ddla_handle_t.h>
#include <ddla/gemmVbatched.h>
#include <ddla/ddla.h>

#elif defined(__MAGMA_RI)
#include "Magmablas_Interface-Contiguous.h"
Expand All @@ -28,6 +26,7 @@

#include <mpi.h>

#include <complex>
#include <cstdio>
#include <cstdlib>
#include <stdexcept>
Expand All @@ -41,16 +40,12 @@ namespace GPU_Backend

#ifdef __DDLA_RI
using Int = int;
using Transpose = ddla::deblasOperation_t;
using Transpose = char;
using Queue = ddla::DdlaHandle_t;
constexpr Transpose NoTrans = ddla::DEBLAS_OP_N;
constexpr Transpose Trans = ddla::DEBLAS_OP_T;
constexpr Transpose ConjTrans = ddla::DEBLAS_OP_C;
constexpr Transpose NoTrans = 'N';
constexpr Transpose Trans = 'T';
constexpr Transpose ConjTrans = 'C';

inline ddla::deviceStream_t stream(Queue queue)
{
return static_cast<ddla::deviceStream_t>(ddla::ddla_get_stream(queue));
}
#elif defined(__MAGMA_RI)
using Int = magma_int_t;
using Transpose = magma_trans_t;
Expand Down Expand Up @@ -92,7 +87,7 @@ template<typename T_GPU> Type_to_CPU<T_GPU> data_to_CPU(const T_GPU data_gpu)
#endif


inline void validate_local_gpu_count(const MPI_Comm& communicator)
inline void validate_local_gpu_count(const MPI_Comm& communicator, int device_count)
{
MPI_Comm local_communicator = MPI_COMM_NULL;
if (MPI_Comm_split_type(
Expand All @@ -105,12 +100,6 @@ inline void validate_local_gpu_count(const MPI_Comm& communicator)
MPI_Comm_size(local_communicator, &local_size);
MPI_Comm_free(&local_communicator);

int device_count = 0;
#ifdef __DDLA_RI
ddla::DEVICE_CHECK(ddla::deviceGetDeviceCount(&device_count));
#elif defined(__MAGMA_RI)
device_count = Magma_Wrapper::magma_get_size();
#endif
if (device_count <= 0 || local_size > device_count)
throw std::runtime_error(
"LibRI requires at least one GPU per local MPI rank");
Expand All @@ -122,14 +111,24 @@ class Context
explicit Context(const MPI_Comm& communicator)
{
#ifdef __DDLA_RI
validate_local_gpu_count(communicator);
ddla::ddla_init(queue_);
ddla::ddla_set(queue_, communicator);
ddla::DDLA_CHECK(ddla::ddlaInit(queue_, ddla::DdlaBackend::GPU));
try
{
int device_count = 0;
ddla::DDLA_CHECK(ddla::ddlaGetDeviceCount(queue_, &device_count));
validate_local_gpu_count(communicator, device_count);
ddla::DDLA_CHECK(ddla::ddlaSet(queue_, communicator));
}
catch (...)
{
ddla::ddlaDestroy(queue_);
throw;
}
#elif defined(__MAGMA_RI)
MAGMA_CHECK_RI(magma_init());
try
{
validate_local_gpu_count(communicator);
validate_local_gpu_count(communicator, Magma_Wrapper::magma_get_size());
}
catch (...)
{
Expand All @@ -149,7 +148,7 @@ class Context
~Context()
{
#ifdef __DDLA_RI
ddla::ddla_destroy(queue_);
ddla::ddlaDestroy(queue_);
#elif defined(__MAGMA_RI)
if (queue_ != nullptr)
magma_queue_destroy(queue_);
Expand All @@ -166,28 +165,30 @@ class Context
};

template <typename T>
inline void allocate(T** pointer, std::size_t count)
inline void allocate(T** pointer, std::size_t count, Queue queue)
{
#ifdef __DDLA_RI
ddla::DDLA_CHECK(ddla::ddlaMalloc(queue, pointer, count * sizeof(T)));
#elif defined(__MAGMA_RI)
(void)queue;
if (count == 0)
{
*pointer = nullptr;
return;
}
#ifdef __DDLA_RI
ddla::DEVICE_CHECK(ddla::deviceMalloc(pointer, count * sizeof(T)));
#elif defined(__MAGMA_RI)
MAGMA_CHECK_RI(magma_malloc(reinterpret_cast<void**>(pointer), count * sizeof(T)));
#endif
}

template <typename T>
inline void free(T* pointer)
inline void free(T* pointer, Queue queue)
{
if (pointer == nullptr)
return;
#ifdef __DDLA_RI
ddla::DEVICE_CHECK(ddla::deviceFree(pointer));
ddla::DDLA_CHECK(ddla::ddlaFree(queue, pointer));
#elif defined(__MAGMA_RI)
(void)queue;
MAGMA_CHECK_RI(magma_free(pointer));
#endif
}
Expand All @@ -199,9 +200,8 @@ inline void upload(
if (count == 0)
return;
#ifdef __DDLA_RI
ddla::DEVICE_CHECK(deviceMemcpyAsync(
device_data, host_data, count * sizeof(T),
ddla::deviceMemcpyHostToDevice, stream(queue)));
ddla::DDLA_CHECK(ddla::ddlaMemcpy(
queue, device_data, host_data, count * sizeof(T)));
#elif defined(__MAGMA_RI)
magma_setvector_async(
count, sizeof(T), host_data, 1, device_data, 1, queue);
Expand All @@ -215,9 +215,8 @@ inline void download(
if (count == 0)
return;
#ifdef __DDLA_RI
ddla::DEVICE_CHECK(deviceMemcpyAsync(
host_data, device_data, count * sizeof(T),
ddla::deviceMemcpyDeviceToHost, stream(queue)));
ddla::DDLA_CHECK(ddla::ddlaMemcpy(
queue, host_data, device_data, count * sizeof(T)));
#elif defined(__MAGMA_RI)
magma_getvector_async(
count, sizeof(T), device_data, 1, host_data, 1, queue);
Expand All @@ -229,8 +228,7 @@ inline void memset(void* pointer, int value, std::size_t bytes, Queue queue)
if (bytes == 0)
return;
#ifdef __DDLA_RI
ddla::DEVICE_CHECK(deviceMemsetAsync(
pointer, value, bytes, stream(queue)));
ddla::DDLA_CHECK(ddla::ddlaMemset(queue, pointer, value, bytes));
#elif defined(__MAGMA_RI)
MAGMA_CHECK_RI(magma_memset(pointer, value, bytes));
#endif
Expand All @@ -239,8 +237,7 @@ inline void memset(void* pointer, int value, std::size_t bytes, Queue queue)
inline void sync(Queue queue)
{
#ifdef __DDLA_RI
ddla::DEVICE_CHECK(
ddla::deviceStreamSynchronize(stream(queue)));
ddla::DDLA_CHECK(ddla::ddlaSynchronize(queue));
#elif defined(__MAGMA_RI)
magma_queue_sync(queue);
#endif
Expand All @@ -257,10 +254,10 @@ inline void gemmVbatched(
Int*& lda = transA == NoTrans ? k : m;
Int*& ldb = transB == NoTrans ? n : k;
Int*& ldc = n;
ddla::gemmVbatched(
transB, transA, n, m, k,
ddla::gemmVbatched<ddla::DdlaBackend::GPU>(
queue, transB, transA, n, m, k,
alpha, dB_array, ldb, dA_array, lda,
beta, dC_array, ldc, batch_count, queue);
beta, dC_array, ldc, batch_count);
#elif defined(__MAGMA_RI)
Magma_Interface::gemm_vbatched(
transA, transB, m, n, k,
Expand Down Expand Up @@ -290,14 +287,14 @@ inline void gemmVbatched2s(
Int* lda_1 = transA_1 == NoTrans ? k_1 : m_1;
Int* ldb_1 = transB_1 == NoTrans ? n_1 : k_1;
Int* ldc_1 = n_1;
ddla::gemmVbatched2s(
transB_0, transA_0, n_0, m_0, k_0,
ddla::gemmVbatched2s<ddla::DdlaBackend::GPU>(
queue, transB_0, transA_0, n_0, m_0, k_0,
alpha_0, dB_array_0, ldb_0, dA_array_0, lda_0,
beta_0, dC_array_0, ldc_0,
transB_1, transA_1, n_1, m_1, k_1,
alpha_1, dAB_array_1, ldb_1, lda_1,
beta_1, dC_array_1, ldc_1, !C0_left,
batch_count, segment_sizes.data(), segment_sizes.size(), queue);
batch_count, segment_sizes.data(), segment_sizes.size());
#elif defined(__MAGMA_RI)
Magma_Interface::gemm_vbatched_2s(
transA_0, transB_0, m_0, n_0, k_0,
Expand Down
6 changes: 4 additions & 2 deletions include/RI/ri/gpu/Dim.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ namespace RI
class Dim
{
public:
GPU_Backend::Queue queue_ = nullptr;
void input(const int shape)
{
this->shapes.push_back(shape);
Expand All @@ -27,14 +28,15 @@ class Dim

void upload(GPU_Backend::Queue queue)
{
GPU_Backend::allocate(&this->d_m, this->shapes.size() + 1);
this->queue_ = queue;
GPU_Backend::allocate(&this->d_m, this->shapes.size() + 1, queue);
GPU_Backend::upload(
this->shapes.size(), this->shapes.data(), this->d_m, queue);
}

~Dim()
{
GPU_Backend::free(this->d_m);
GPU_Backend::free(this->d_m, this->queue_);
}

int* data() const
Expand Down
10 changes: 6 additions & 4 deletions include/RI/ri/gpu/GPU_Data_Input.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ template<typename TA, typename TAC, typename Tdata_CPU>
class Input
{
public:
GPU_Backend::Queue queue_ = nullptr;
using Tdata_GPU = GPU_Backend::Type_to_GPU<Tdata_CPU>;
using Same_data_layout = std::integral_constant<
bool,
Expand Down Expand Up @@ -63,11 +64,12 @@ class Input
// 将对应的 C, V, D 放在 d_Cs, d_Vs, d_Ds 上
void upload(GPU_Backend::Queue queue)
{
this->queue_ = queue;
std::vector<std::size_t> h_data_begin(this->h_data.size()+1, 0);
for(std::size_t i=1; i<h_data_begin.size(); ++i)
h_data_begin[i] = h_data_begin[i-1] + this->h_data[i-1].size();

GPU_Backend::allocate(&this->d_data, h_data_begin.back());
GPU_Backend::allocate(&this->d_data, h_data_begin.back(), queue);
for(std::size_t i=0; i<this->h_data.size(); ++i)
GPU_Backend::upload(
this->h_data[i].size(), this->h_data[i].data(),
Expand All @@ -77,14 +79,14 @@ class Input
std::vector<Tdata_GPU*> d_array_(batchCount); // 记录每个batch的 d_data 指针（CPU）
for(std::size_t i=0; i<batchCount; ++i)
d_array_[i] = this->d_data + this->h_array[i].pos + h_data_begin[this->h_array[i].thread_num];
GPU_Backend::allocate(&this->d_array, batchCount);
GPU_Backend::allocate(&this->d_array, batchCount, queue);
GPU_Backend::upload(batchCount, d_array_.data(), this->d_array, queue);
}

~Input()
{
GPU_Backend::free(this->d_data);
GPU_Backend::free(this->d_array);
GPU_Backend::free(this->d_data, this->queue_);
GPU_Backend::free(this->d_array, this->queue_);
}

private:
Expand Down
15 changes: 9 additions & 6 deletions include/RI/ri/gpu/GPU_Data_Mul.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ template<typename TA, typename TAC, typename Tdata_CPU>
class Mul
{
public:
GPU_Backend::Queue queue_ = nullptr;
using Tdata_GPU = GPU_Backend::Type_to_GPU<Tdata_CPU>;

const Pack &insert_1st(const TA &Aa, const TAC &Ab, const Shape_Vector &shape)
Expand All @@ -38,15 +39,16 @@ class Mul

void upload_1st(GPU_Backend::Queue queue)
{
GPU_Backend::allocate(&this->d_data, this->totalSize);
this->queue_ = queue;
GPU_Backend::allocate(&this->d_data, this->totalSize, queue);
GPU_Backend::memset(
this->d_data, 0, totalSize * sizeof(Tdata_GPU), queue);

const std::size_t batchCount = this->h_array_1.size();
std::vector<Tdata_GPU*> d_array_(batchCount); // 记录每个batch的 d_data 指针（CPU）
for (std::size_t i = 0; i < batchCount; i++)
d_array_[i] = this->d_data + this->h_array_1[i].pos;
GPU_Backend::allocate(&this->d_array_1, batchCount);
GPU_Backend::allocate(&this->d_array_1, batchCount, queue);
GPU_Backend::upload(
batchCount, d_array_.data(), this->d_array_1, queue);
}
Expand All @@ -65,20 +67,21 @@ class Mul
// 在 GPU 上分配空间，并获取指针
void upload_2nd(GPU_Backend::Queue queue)
{
this->queue_ = queue;
const std::size_t batchCount = this->h_array_2.size();
std::vector<Tdata_GPU*> d_array_2_(batchCount); // 记录每个batch的 d_data 指针（CPU）
for (std::size_t i = 0; i < batchCount; i++)
d_array_2_[i] = this->d_data + this->h_array_2[i].pos;
GPU_Backend::allocate(&this->d_array_2, batchCount);
GPU_Backend::allocate(&this->d_array_2, batchCount, queue);
GPU_Backend::upload(
batchCount, d_array_2_.data(), this->d_array_2, queue);
}

~Mul()
{
GPU_Backend::free(this->d_data);
GPU_Backend::free(this->d_array_1);
GPU_Backend::free(this->d_array_2);
GPU_Backend::free(this->d_data, this->queue_);
GPU_Backend::free(this->d_array_1, this->queue_);
GPU_Backend::free(this->d_array_2, this->queue_);
}

std::size_t totalSize = 0; // 总的数据数量
Expand Down
10 changes: 6 additions & 4 deletions include/RI/ri/gpu/GPU_Data_Output.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ template<typename TA, typename TAC, typename Tdata_CPU>
class Output
{
public:
GPU_Backend::Queue queue_ = nullptr;
using Tdata_GPU = GPU_Backend::Type_to_GPU<Tdata_CPU>;
using Same_data_layout = std::integral_constant<
bool,
Expand All @@ -46,15 +47,16 @@ class Output

void upload(GPU_Backend::Queue queue)
{
GPU_Backend::allocate(&this->d_data, this->totalSize);
this->queue_ = queue;
GPU_Backend::allocate(&this->d_data, this->totalSize, queue);
GPU_Backend::memset(
this->d_data, 0, totalSize * sizeof(Tdata_GPU), queue);

const std::size_t batchCount = this->h_array.size();
std::vector<Tdata_GPU*> d_array_(batchCount); // 记录每个batch的 d_data 指针（CPU）
for (std::size_t i = 0; i < batchCount; i++)
d_array_[i] = this->d_data + this->h_array[i].pos;
GPU_Backend::allocate(&this->d_array, batchCount);
GPU_Backend::allocate(&this->d_array, batchCount, queue);
GPU_Backend::upload(batchCount, d_array_.data(), this->d_array, queue);
}

Expand Down Expand Up @@ -83,8 +85,8 @@ class Output

~Output()
{
GPU_Backend::free(this->d_data);
GPU_Backend::free(this->d_array);
GPU_Backend::free(this->d_data, this->queue_);
GPU_Backend::free(this->d_array, this->queue_);
}

private:
Expand Down
Loading