mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-14 16:33:59 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
101 lines
4.0 KiB
Python
101 lines
4.0 KiB
Python
#!/usr/bin/env python3
|
|
# GMMU=0 MLX_PCI=0000:41:00.0 PYTHONPATH=. python3 extra/mlx_driver/loopback.py
|
|
import struct
|
|
from tinygrad.helpers import getenv, round_up
|
|
from tinygrad.device import Device, BufferSpec, Buffer
|
|
from tinygrad.dtype import dtypes
|
|
from tinygrad.runtime.support.system import PCIDevice
|
|
from tinygrad.runtime.support.memory import AddrSpace
|
|
from tinygrad.runtime.ops_amd import AMDComputeQueue
|
|
from tinygrad.helpers import to_be32, to_be64
|
|
from extra.mlx_driver.mlxdev import MLXDev, MLXQP
|
|
|
|
BUF_SIZE = 0x1000
|
|
MLX_PCI = getenv("MLX_PCI", "0000:41:00.0")
|
|
MLX_IP = getenv("MLX_IP", "10.0.0.1")
|
|
|
|
def map_phys_to_gpu(gpu, paddr, size):
|
|
size = round_up(size, 0x1000)
|
|
va = gpu.iface.dev_impl.mm.alloc_vaddr(size, align=0x1000)
|
|
gpu.iface.dev_impl.mm.map_range(va, size, [(paddr, size)], aspace=AddrSpace.SYS, snooped=True, uncached=True)
|
|
return va
|
|
|
|
print("[init] AMD GPU...")
|
|
gpu = Device["AMD"]
|
|
|
|
print(f"[init] MLX5 at {MLX_PCI}")
|
|
dev = MLXDev(PCIDevice("mlx5", MLX_PCI), ip=MLX_IP)
|
|
qp = MLXQP(dev)
|
|
|
|
print(f"[init] loopback connect QP 0x{qp.qp_info['qpn']:x}")
|
|
qp.connect(qp.qp_info['qpn'], dev.mac, int.from_bytes(dev.local_gid, 'big'))
|
|
|
|
# allocate src/dst via AMD GPU allocator
|
|
buf_src = Buffer(gpu.device, BUF_SIZE, dtypes.uint8, options=BufferSpec(nolru=True), preallocate=True)
|
|
buf_dst = Buffer(gpu.device, BUF_SIZE, dtypes.uint8, options=BufferSpec(nolru=True), preallocate=True)
|
|
|
|
bar_base = gpu.iface.pci_dev.bar_info(gpu.iface.vram_bar)[0]
|
|
src_paddr = buf_src.meta.mapping.paddrs[0][0] + bar_base
|
|
dst_paddr = buf_dst.meta.mapping.paddrs[0][0] + bar_base
|
|
print(f"src paddr=0x{src_paddr:x} dst paddr=0x{dst_paddr:x}")
|
|
|
|
# fill src, zero dst
|
|
test_msg = b"Hello from loopback send/recv!"
|
|
buf_src.copy_from(Buffer("PYTHON", BUF_SIZE, dtypes.uint8, opaque=memoryview(bytearray(test_msg.ljust(BUF_SIZE, b'\x00')))))
|
|
buf_dst.copy_from(Buffer("PYTHON", BUF_SIZE, dtypes.uint8, opaque=memoryview(bytearray(BUF_SIZE))))
|
|
gpu.synchronize()
|
|
|
|
# post recv WQE on RQ from CPU (scatter entry: byte_count, lkey, addr)
|
|
rq_mask = (1 << 4) - 1 # log_rq_size=4
|
|
rq_wqe = qp.qp_buf.view((qp.rq_head & rq_mask) * 16, 16)
|
|
rq_wqe[:] = struct.pack('>IIQ', len(test_msg), dev.mkey, dst_paddr)
|
|
qp.rq_head += 1
|
|
# ring recv doorbell from CPU (DBR offset 0 = recv counter)
|
|
dev.dbr[qp.qp_dbr // 4] = to_be32(qp.rq_head)
|
|
|
|
# build send WQE in SQ from CPU (opcode 0x0a = SEND, ds_count=2)
|
|
sq_head = qp.sq_head
|
|
sq_mask = (1 << qp.log_sq_size) - 1
|
|
wqe = qp.qp_buf.view(qp.sq_offset + (sq_head & sq_mask) * 64, 64)
|
|
wqe[:] = bytes(64)
|
|
wqe[0:8] = struct.pack('>II', (sq_head << 8) | 0x0a, (qp.qp_info['qpn'] << 8) | 2)
|
|
wqe[11] = 0x08 # CE: signal completion
|
|
wqe[16:32] = struct.pack('>IIQ', len(test_msg), dev.mkey, src_paddr)
|
|
qp.sq_head += 1
|
|
doorbell_val = to_be64(int.from_bytes(bytes(wqe[0:8]), 'big'))
|
|
|
|
# map MLX5 UAR and DBR into GPU VA
|
|
uar_paddr = dev.pci_dev.bar_info(0)[0] + dev.uar * 0x1000
|
|
uar_gpu_va = map_phys_to_gpu(gpu, uar_paddr, 0x1000)
|
|
dbr_gpu_va = map_phys_to_gpu(gpu, dev.dbr_paddrs[0], 0x1000)
|
|
print(f"UAR gpu_va=0x{uar_gpu_va:x} DBR gpu_va=0x{dbr_gpu_va:x}")
|
|
|
|
# GPU rings send doorbell via compute queue release_mem
|
|
q = AMDComputeQueue(gpu)
|
|
q.wait(gpu.timeline_signal, gpu.timeline_value - 1)
|
|
# write DBR (32-bit sq_head) - send doorbell at qp_dbr + 4
|
|
q.release_mem(dbr_gpu_va + qp.qp_dbr + 4, to_be32(qp.sq_head), q.pm4.data_sel__mec_release_mem__send_32_bit_low,
|
|
q.pm4.int_sel__mec_release_mem__none)
|
|
# write UAR doorbell (64-bit)
|
|
q.release_mem(uar_gpu_va + 0x800, doorbell_val, q.pm4.data_sel__mec_release_mem__send_64_bit_data,
|
|
q.pm4.int_sel__mec_release_mem__none)
|
|
q.signal(gpu.timeline_signal, gpu.next_timeline())
|
|
q.submit(gpu)
|
|
|
|
print("GPU kicked doorbell, waiting...")
|
|
gpu.synchronize()
|
|
|
|
# poll CQ from CPU (send + recv completions)
|
|
qp.poll_cq()
|
|
qp.poll_cq()
|
|
|
|
# read back
|
|
result = bytearray(BUF_SIZE)
|
|
result[:] = buf_dst.as_memoryview()
|
|
gpu.synchronize()
|
|
|
|
got = bytes(result[:len(test_msg)])
|
|
print(f"result: {got}")
|
|
assert got == test_msg, f"MISMATCH: {got} != {test_msg}"
|
|
print("RDMA loopback send/recv test passed (GPU-kicked)")
|