Ví dụ về PJRT

Ví dụ: Trình bổ trợ CUDA JAX

1. Triển khai API PJRT C thông qua trình bao bọc (pjrt_c_api_gpu.h).
2. Thiết lập điểm truy cập cho gói (setup.py).
3. Triển khai phương thức initialize() (__init__.py).
4. Có thể được kiểm thử bằng bất kỳ kiểm thử jax nào cho CUDA.

Triển khai khung

Một số tài liệu tham khảo về cách sử dụng PJRT ở phía khung để giao tiếp với các thiết bị PJRT:

• JAX
  • jax-ml/jax tương tác với các API PJRT thông qua các API xla_client
• GoMLX
  • gomlx/gopjrt
  • gomlx/gomlx > backends/xla
• ZML
  • Trình bao bọc API PJRT pjrt.zig
  • Tải trình bổ trợ PJRT context.zig
  • Tương tác với vùng đệm PJRT buffer.zig
  • Thực thi mô-đun thông qua PJRT module.zig

Triển khai phần cứng

• Trình bổ trợ tích hợp đầy đủ (PJRT+MLIR+XLA):
  • Trình bổ trợ CPU XLA
  • Trình bổ trợ GPU XLA
  • Trình bổ trợ Intel XLA
• Trình bổ trợ tích hợp nhẹ (PJRT+MLIR):
  • Trình bổ trợ Trình diễn giải tham chiếu StableHLO (trình bổ trợ C++, dựa trên MLIR, sẽ được liên kết sau các lớp học lập trình)
  • Trình bổ trợ Tenstorrent-XLA (trình bổ trợ C dựa trên MLIR)