ตัวอย่าง PJRT

ตัวอย่าง: ปลั๊กอิน JAX CUDA

1. การใช้งาน PJRT C API ผ่าน Wrapper (pjrt_c_api_gpu.h)
2. ตั้งค่าจุดแรกเข้าของแพ็กเกจ (setup.py)
3. ใช้เมธอด initialize() (__init__.py)
4. ทดสอบได้โดยใช้การทดสอบ jax สำหรับ CUDA

การติดตั้งใช้งานเฟรมเวิร์ก

ข้อมูลอ้างอิงบางส่วนสำหรับการใช้ PJRT ในด้านเฟรมเวิร์กเพื่อเชื่อมต่อกับอุปกรณ์ PJRT มีดังนี้

• JAX
  • jax-ml/jax โต้ตอบกับ PJRT API ผ่าน xla_client API
• GoMLX
  • gomlx/gopjrt
  • gomlx/gomlx > backends/xla
• ZML
  • Wrapper ของ PJRT API pjrt.zig
  • โหลดปลั๊กอิน PJRT context.zig
  • การโต้ตอบกับบัฟเฟอร์ PJRT buffer.zig
  • เรียกใช้โมดูลผ่าน PJRT module.zig

การติดตั้งใช้งานฮาร์ดแวร์

• ปลั๊กอินการผสานรวมแบบสมบูรณ์ (PJRT+MLIR+XLA):
  • ปลั๊กอิน CPU สำหรับ XLA
  • ปลั๊กอิน GPU สำหรับ XLA
  • ปลั๊กอิน Intel XLA
• ปลั๊กอินการผสานรวมแบบเบา (PJRT+MLIR):
  • ปลั๊กอินโปรแกรมแปลภาษาอ้างอิง StableHLO (ปลั๊กอิน C++ ที่อิงตาม MLIR ซึ่งจะลิงก์หลังจาก devlabs)
  • ปลั๊กอิน Tenstorrent-XLA (ปลั๊กอิน C ที่ใช้ MLIR)