系统软件层

OpenXLA 通过使用 CUDA Driver API 与底层 GPU 进行交互，主要流程涉及从高层的模型执行到底层 GPU 资源的管理和调度：

  • OpenXLA 接收从 TensorFlow、PyTorch 等深度学习框架导出的模型（如 StableHLO 或 ONNX 格式）。

  • OpenXLA Runtime 解析模型，并通过图优化技术，减少冗余计算、调整执行顺序，以提升运行效率。

  • OpenXLA 会将优化后的模型图分解为多个子任务（操作），并对这些操作进行调度，分配给适合的硬件资源（如 CPU、GPU、TPU 等）。

  • 通过 CUDA Driver API，OpenXLA 初始化 GPU 设备，包括检测可用的 GPU 设备、分配 GPU 资源等。

  • OpenXLA 通过 CUDA Driver API 请求 GPU 上的内存分配。使用 cuMemAlloc 等 API 在 GPU 上分配全局内存。通过 cuMemcpyHtoD（Host to Device）等 API 将 CPU 上的数据传输到 GPU 内存中。

  • OpenXLA 通过 CUDA Driver API 提交并启动计算核函数（Kernel），通常使用 cuLaunchKernel 启动并行计算。

  • 核函数会在 GPU 的并行计算单元（CUDA 核心）上运行，处理计算任务。

  • OpenXLA 负责为这些核函数设置执行配置（如网格大小、线程块大小等），确保高效并行执行。

  • 在核函数执行结束后，OpenXLA 使用 CUDA Driver API 来同步执行流程，确保 GPU 的计算结果准备就绪，典型的 API 有 cuCtxSynchronize。

  • 同时，OpenXLA 监控 GPU 执行的状态，处理可能的错误（如内存不足、非法访问等）。

  • 通过 cuMemcpyDtoH（Device to Host）将 GPU 计算结果传回到 CPU 进行进一步处理。

  • 完成计算后，OpenXLA 通过 cuMemFree 等 API 释放在 GPU 上分配的内存资源。