IQ.Pilot Prebuilt Release @ ab07000

tinygrad_repo/extra/mlx_driver/connect.py

@@ -0,0 +1,68 @@
+#!/usr/bin/env python3
+import subprocess, json, sys, os
+
+REMOTE_HOST = os.getenv("REMOTE_HOST", "192.168.52.154")
+LOCAL_PCI   = os.getenv("MLX_PCI", "0000:41:00.0")
+REMOTE_PCI  = os.getenv("REMOTE_PCI", "0000:41:00.0")
+LOCAL_IP    = os.getenv("LOCAL_IP", "10.0.0.1")
+REMOTE_IP   = os.getenv("REMOTE_IP", "10.0.0.2")
+SSH         = ["ssh", "-o", "StrictHostKeyChecking=no", REMOTE_HOST]
+TINYGRAD    = os.path.dirname(os.path.abspath(__file__)) + "/../.."
+
+print("syncing code to remote")
+subprocess.run(["rsync", "-az", "--exclude=.git", "--exclude=__pycache__", "--exclude=*.pyc",
+                TINYGRAD + "/", f"{REMOTE_HOST}:~/tinygrad/"], check=True)
+
+print("booting remote")
+remote = subprocess.Popen(
+  SSH + [f"cd ~/tinygrad && sudo PYTHONPATH=. MLX_DEBUG=1 MLX_PCI={REMOTE_PCI} MLX_IP={REMOTE_IP} python3 extra/mlx_driver/mlxdev.py --server"],
+  stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=sys.stderr, text=True)
+
+remote_info = None
+for line in iter(remote.stdout.readline, ''):
+  print(f"  [remote] {line}", end='')
+  try: remote_info = json.loads(line.strip()); break
+  except json.JSONDecodeError: pass
+assert remote_info, "failed to get remote connection info"
+
+print("booting local")
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "../.."))
+from extra.mlx_driver.mlxdev import MLXDev, MLXQP
+from tinygrad.runtime.support.system import PCIDevice
+
+local_dev = MLXDev(PCIDevice("mlx5", LOCAL_PCI), ip=LOCAL_IP)
+local_qp = MLXQP(local_dev)
+local_info = {"qpn": local_qp.qpn, "mac": local_dev.mac.to_bytes(6,'big').hex(), "gid": local_dev.local_gid.hex()}
+
+remote.stdin.write(json.dumps(local_info) + "\n")
+remote.stdin.flush()
+for line in iter(remote.stdout.readline, ''):
+  print(f"  [remote] {line}", end='')
+  if "connected" in line: break
+
+local_qp.connect(remote_info["qpn"], int(remote_info["mac"], 16), int(remote_info["gid"], 16))
+print("both QPs in RTS")
+
+remote_target = None
+for line in iter(remote.stdout.readline, ''):
+  print(f"  [remote] {line}", end='')
+  try: remote_target = json.loads(line.strip()); break
+  except json.JSONDecodeError: pass
+assert remote_target
+
+test_msg = b"Test message, rdma works!"
+src_mem, src_paddrs = local_dev.pci_dev.alloc_sysmem(0x1000)
+for i, b in enumerate(test_msg): src_mem[i] = b
+
+print(f"RDMA WRITE {len(test_msg)}B to remote phys 0x{remote_target['target_addr']:x}")
+local_qp.rdma_write(remote_target["target_addr"], remote_target["rkey"], src_paddrs[0], local_dev.mkey, len(test_msg))
+
+remote.stdin.write("done\n")
+remote.stdin.flush()
+for line in iter(remote.stdout.readline, ''):
+  print(f"  [remote] {line}", end='')
+  if "AS TEXT" in line: break
+
+remote.stdin.close()
+remote.wait()
+print("RDMA WRITE test complete")

tinygrad_repo/extra/mlx_driver/loopback.py

@@ -0,0 +1,99 @@
+#!/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
+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 = gpu.allocator.alloc(BUF_SIZE, BufferSpec(nolru=True))
+buf_dst = gpu.allocator.alloc(BUF_SIZE, BufferSpec(nolru=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!"
+gpu.allocator._copyin(buf_src, memoryview(bytearray(test_msg.ljust(BUF_SIZE, b'\x00'))))
+gpu.allocator._copyin(buf_dst, 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)
+gpu.allocator._copyout(memoryview(result), buf_dst)
+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)")