Skip to main content

MUSA Driver API 指南

MUSA Driver API 提供底层编程接口,用于显式上下文管理、模块加载、高级设备控制。

  • 设计特点:显式控制(需要显式初始化和创建上下文)、灵活性高(可以精细控制设备资源)
  • 适用场景:多 GPU 应用、多进程应用、库开发、需要精细控制的高级应用

概述

Runtime API vs Driver API 对比

特性Runtime APIDriver API
初始化隐式(首次调用自动初始化)显式(muInit()
上下文管理自动(primary context)手动创建和销毁
代码复杂度
灵活性
适用场景应用程序、快速原型多 GPU、多进程、库开发

对比 Runtime API:详见 Runtime API 指南


初始化和上下文管理

显式初始化

#include <mu.h>

// 显式初始化(可选,首次调用 API 时也会自动初始化)
muInit(0);

// 获取设备
MUdevice device;
muDeviceGet(&device, 0);

// 创建上下文
MUcontext context;
muCtxCreate(&context, 0, device);

// 设置当前上下文
muCtxSetCurrent(context);

// 使用完毕后销毁上下文
muCtxDestroy(context);

上下文管理

// 获取当前上下文
MUcontext currentCtx;
muCtxGetCurrent(&currentCtx);

// 获取上下文的设备
MUdevice device;
muCtxGetDevice(&device, context);

// 获取上下文标志
unsigned int flags;
muCtxGetFlags(&flags, context);

// 切换上下文
muCtxSetCurrent(context1);
// ... 使用 context1 ...
muCtxSetCurrent(context2);
// ... 使用 context2 ...

主上下文

// 获取设备的 primary context(Runtime API 使用)
MUcontext primaryCtx;
muDevicePrimaryCtxGetState(device, &primaryCtx, NULL);

// 释放 primary context(通常不需要手动调用)
muDevicePrimaryCtxRelease(device);

模块管理

模块加载

// 从 PTX 文件加载模块
MUmodule module;
muModuleLoad(&module, "kernel.ptx");

// 从 fatbinary 加载模块
muModuleLoad(&module, "kernel.fatbin");

// 从数据加载模块
void* moduleData = ...;
muModuleLoadData(&module, moduleData);

// 卸载模块
muModuleUnload(module);

获取内核函数

// 从模块中获取内核函数
MUfunction kernelFunc;
muModuleGetFunction(&kernelFunc, module, "vectorAdd");

// 获取内核参数大小
int paramSize;
muFuncGetParamSize(kernelFunc, &paramSize);

// 设置内核参数
muParamSeti(kernelFunc, 0, n);
muParamSetd(kernelFunc, sizeof(int), d_a);
muParamSetd(kernelFunc, sizeof(int) * 2, d_b);
muParamSetd(kernelFunc, sizeof(int) * 3, d_c);

启动内核

// 设置块大小
muParamSetSize(kernelFunc, blockSize);

// 启动内核
muLaunchKernel(kernelFunc, 1, 1, 1, gridSize, 1, 1, 0, stream);

// 或使用简化 API
muLaunchKernelByPtr(kernelFunc, gridSize, blockSize, 0, stream);

内存管理

设备内存分配

// 分配设备内存
MUdeviceptr d_data;
muMemAlloc(&d_data, size);

// 分配 2D 内存
MUdeviceptr d_2d;
size_t pitch;
muMemAllocPitch(&d_2d, &pitch, widthInBytes, height, 4);

// 释放内存
muMemFree(d_data);

主机内存分配

// 分配主机内存
void* h_data;
muMemAllocHost(&h_data, size);

// 分配写合并内存
void* h_wc;
muMemHostAlloc(&h_wc, size, MU_MEMHOSTALLOC_WRITECOMBINED);

// 注册已存在的主机内存
muMemHostRegister(h_existing, size, MU_MEMHOSTREGISTER_DEVICEMAP);

// 获取设备指针
MUdeviceptr d_ptr;
muMemHostGetDevicePointer(&d_ptr, h_data, 0);

// 释放主机内存
muMemFreeHost(h_data);
muMemHostUnregister(h_existing);

统一内存

// 分配统一内存
void* managed;
muMemAllocManaged(&managed, size, MU_MEM_ATTACH_GLOBAL);

// 内存建议
muMemAdvise(managed, size, MU_MEM_ADVISE_SET_PREFERRED_LOCATION, deviceId);

// 内存预取
muMemPrefetchAsync(managed, size, deviceId, stream);

内存拷贝

// Host → Device
muMemcpyDtoH(h_dst, d_src, size);

// Device → Host
muMemcpyHtoD(d_dst, h_src, size);

// Device → Device
muMemcpyDtoD(d_dst, d_src, size);

// 2D 内存拷贝
muMemcpy2D(dstPtr, dstPitch, srcPtr, srcPitch, width, height, copyKind);

// 异步拷贝
muMemcpyAsync(d_dst, h_src, size, stream);

流和事件

流管理

// 创建流
MUstream stream;
muStreamCreate(&stream);

// 创建带优先级的流
muStreamCreateWithPriority(&stream, 0, minPriority);

// 创建非阻塞流
muStreamCreateWithPriority(&stream, MU_STREAM_NON_BLOCKING, 0);

// 销毁流
muStreamDestroy(stream);

// 同步流
muStreamSynchronize(stream);

// 查询流状态
muError_t err = muStreamQuery(stream);

事件管理

// 创建事件
MUevent event;
muEventCreate(&event);

// 创建带标志的事件
muEventCreateWithFlags(&event, MU_EVENT_DISABLE_TIMING);

// 记录事件
muEventRecord(event, stream);

// 等待事件完成
muEventSynchronize(event);

// 查询事件状态
muError_t err = muEventQuery(event);

// 计算事件间隔
float elapsedMs;
muEventElapsedTime(&elapsedMs, startEvent, stopEvent);

// 销毁事件
muEventDestroy(event);

设备管理

设备查询

// 获取设备数量
int count;
muGetDeviceCount(&count);

// 获取设备
MUdevice device;
muDeviceGet(&device, 0);

// 获取设备属性
int value;
muDeviceGetAttribute(&value, MU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK, device);

// 获取设备名称
char name[256];
muDeviceGetName(name, sizeof(name), device);

// 获取设备 UUID
MUuuid uuid;
muDeviceGetUuid(&uuid, device);

设备属性

// 常用属性查询
int maxThreadsPerBlock;
muDeviceGetAttribute(&maxThreadsPerBlock,
MU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK, device);

int sharedMemPerBlock;
muDeviceGetAttribute(&sharedMemPerBlock,
MU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK, device);

int multiProcessorCount;
muDeviceGetAttribute(&multiProcessorCount,
MU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, device);

int computeCapabilityMajor, computeCapabilityMinor;
muDeviceGetAttribute(&computeCapabilityMajor,
MU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, device);
muDeviceGetAttribute(&computeCapabilityMinor,
MU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, device);

错误处理

// 获取错误信息
muError_t err = muGetLastError();
const char* errorStr;
muGetErrorString(&errorStr, err);
printf("Error: %s\n", errorStr);

// 获取错误名称
const char* errorName;
muGetErrorName(&errorName, err);
printf("Error name: %s\n", errorName);

// 错误检查宏
#define MUSA_DRIVER_CHECK(call) \
do { \
muError_t err = call; \
if (err != muSuccess) { \
const char* errorStr; \
muGetErrorString(&errorStr, err); \
fprintf(stderr, "Driver API error at %s:%d: %s\n", \
__FILE__, __LINE__, errorStr); \
exit(EXIT_FAILURE); \
} \
} while(0)

完整示例

基础示例

#include <mu.h>
#include <stdio.h>

int main() {
muError_t err;
MUdevice device;
MUcontext context;

// 1. 显式初始化
err = muInit(0);
if (err != muSuccess) {
printf("muInit failed: %d\n", err);
return -1;
}

// 2. 获取设备
err = muDeviceGet(&device, 0);
if (err != muSuccess) {
printf("muDeviceGet failed: %d\n", err);
return -1;
}

// 3. 创建上下文
err = muCtxCreate(&context, 0, device);
if (err != muSuccess) {
printf("muCtxCreate failed: %d\n", err);
return -1;
}

// 4. 设置当前上下文
err = muCtxSetCurrent(context);
if (err != muSuccess) {
printf("muCtxSetCurrent failed: %d\n", err);
return -1;
}

printf("Context created successfully!\n");

// 5. 使用设备...
// ... 内存分配、内核启动等 ...

// 6. 清理
muCtxDestroy(context);

return 0;
}

多 GPU 上下文管理

#include <mu.h>
#include <stdio.h>

#define NUM_DEVICES 2

int main() {
MUdevice devices[NUM_DEVICES];
MUcontext contexts[NUM_DEVICES];

// 初始化
muInit(0);

// 为每个设备创建上下文
for (int i = 0; i < NUM_DEVICES; i++) {
muDeviceGet(&devices[i], i);
muCtxCreate(&contexts[i], 0, devices[i]);
}

// 在不同设备上执行任务
for (int i = 0; i < NUM_DEVICES; i++) {
muCtxSetCurrent(contexts[i]);
printf("Working on device %d\n", i);
// ... 执行任务 ...
}

// 清理
for (int i = 0; i < NUM_DEVICES; i++) {
muCtxDestroy(contexts[i]);
}

return 0;
}

相关文档