IQ.Pilot Release Commit @ bec7652
This commit is contained in:
IQ.Lvbs history cleanup
@@ -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")
|
||||
@@ -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)")
|
||||
144
artifacts/package_sources/tinygrad/extra/mlx_driver/mlx5.h
Normal file
144
artifacts/package_sources/tinygrad/extra/mlx_driver/mlx5.h
Normal file
@@ -0,0 +1,144 @@
|
||||
// MLX5 autogen header — kernel struct layouts and constants
|
||||
typedef unsigned char __u8;
|
||||
typedef unsigned short __be16;
|
||||
typedef unsigned int __be32;
|
||||
typedef unsigned long long __be64;
|
||||
|
||||
// --- device.h structs ---
|
||||
|
||||
struct mlx5_cmd_layout {
|
||||
__u8 type;
|
||||
__u8 rsvd0[3];
|
||||
__be32 inlen;
|
||||
__be64 in_ptr;
|
||||
__be32 in[4];
|
||||
__be32 out[4];
|
||||
__be64 out_ptr;
|
||||
__be32 outlen;
|
||||
__u8 token;
|
||||
__u8 sig;
|
||||
__u8 rsvd1;
|
||||
__u8 status_own;
|
||||
};
|
||||
|
||||
struct mlx5_cmd_prot_block {
|
||||
__u8 data[512];
|
||||
__u8 rsvd0[48];
|
||||
__be64 next;
|
||||
__be32 block_num;
|
||||
__u8 rsvd1;
|
||||
__u8 token;
|
||||
__u8 ctrl_sig;
|
||||
__u8 sig;
|
||||
};
|
||||
|
||||
struct mlx5_init_seg {
|
||||
__be32 fw_rev;
|
||||
__be32 cmdif_rev_fw_sub;
|
||||
__be32 rsvd0[2];
|
||||
__be32 cmdq_addr_h;
|
||||
__be32 cmdq_addr_l_sz;
|
||||
__be32 cmd_dbell;
|
||||
__be32 rsvd1[120];
|
||||
__be32 initializing;
|
||||
};
|
||||
|
||||
// --- Command opcodes (mlx5_ifc.h) ---
|
||||
#define MLX5_CMD_OP_QUERY_HCA_CAP 0x100
|
||||
#define MLX5_CMD_OP_QUERY_ADAPTER 0x101
|
||||
#define MLX5_CMD_OP_INIT_HCA 0x102
|
||||
#define MLX5_CMD_OP_TEARDOWN_HCA 0x103
|
||||
#define MLX5_CMD_OP_ENABLE_HCA 0x104
|
||||
#define MLX5_CMD_OP_DISABLE_HCA 0x105
|
||||
#define MLX5_CMD_OP_QUERY_PAGES 0x107
|
||||
#define MLX5_CMD_OP_MANAGE_PAGES 0x108
|
||||
#define MLX5_CMD_OP_SET_HCA_CAP 0x109
|
||||
#define MLX5_CMD_OP_QUERY_ISSI 0x10a
|
||||
#define MLX5_CMD_OP_SET_ISSI 0x10b
|
||||
#define MLX5_CMD_OP_SET_DRIVER_VERSION 0x10d
|
||||
#define MLX5_CMD_OP_CREATE_MKEY 0x200
|
||||
#define MLX5_CMD_OP_QUERY_SPECIAL_CONTEXTS 0x203
|
||||
#define MLX5_CMD_OP_CREATE_EQ 0x301
|
||||
#define MLX5_CMD_OP_DESTROY_EQ 0x302
|
||||
#define MLX5_CMD_OP_CREATE_CQ 0x400
|
||||
#define MLX5_CMD_OP_DESTROY_CQ 0x401
|
||||
#define MLX5_CMD_OP_CREATE_QP 0x500
|
||||
#define MLX5_CMD_OP_DESTROY_QP 0x501
|
||||
#define MLX5_CMD_OP_RST2INIT_QP 0x502
|
||||
#define MLX5_CMD_OP_INIT2RTR_QP 0x503
|
||||
#define MLX5_CMD_OP_RTR2RTS_QP 0x504
|
||||
#define MLX5_CMD_OP_QUERY_NIC_VPORT_CONTEXT 0x754
|
||||
#define MLX5_CMD_OP_MODIFY_NIC_VPORT_CONTEXT 0x755
|
||||
#define MLX5_CMD_OP_SET_ROCE_ADDRESS 0x761
|
||||
#define MLX5_CMD_OP_ALLOC_PD 0x800
|
||||
#define MLX5_CMD_OP_ALLOC_UAR 0x802
|
||||
#define MLX5_CMD_OP_ACCESS_REG 0x805
|
||||
#define MLX5_CMD_OP_ALLOC_TRANSPORT_DOMAIN 0x816
|
||||
|
||||
// --- Command status (device.h) ---
|
||||
#define MLX5_CMD_STAT_OK 0x0
|
||||
#define MLX5_CMD_STAT_INT_ERR 0x1
|
||||
#define MLX5_CMD_STAT_BAD_OP_ERR 0x2
|
||||
#define MLX5_CMD_STAT_BAD_PARAM_ERR 0x3
|
||||
#define MLX5_CMD_STAT_BAD_SYS_STATE_ERR 0x4
|
||||
#define MLX5_CMD_STAT_BAD_RES_ERR 0x5
|
||||
#define MLX5_CMD_STAT_RES_BUSY 0x6
|
||||
#define MLX5_CMD_STAT_LIM_ERR 0x8
|
||||
#define MLX5_CMD_STAT_BAD_RES_STATE_ERR 0x9
|
||||
#define MLX5_CMD_STAT_NO_RES_ERR 0xf
|
||||
#define MLX5_CMD_STAT_BAD_INP_LEN_ERR 0x50
|
||||
#define MLX5_CMD_STAT_BAD_OUTP_LEN_ERR 0x51
|
||||
|
||||
// --- HCA cap types ---
|
||||
#define MLX5_CAP_GENERAL 0x0
|
||||
#define MLX5_CAP_ODP 0x2
|
||||
#define MLX5_CAP_ATOMIC 0x3
|
||||
#define MLX5_CAP_ROCE 0x4
|
||||
#define HCA_CAP_OPMOD_GET_MAX 0
|
||||
#define HCA_CAP_OPMOD_GET_CUR 1
|
||||
|
||||
// --- Pages ---
|
||||
#define MLX5_PAGES_GIVE 1
|
||||
#define MLX5_PAGES_TAKE 2
|
||||
#define MLX5_BOOT_PAGES 1
|
||||
#define MLX5_INIT_PAGES 2
|
||||
|
||||
// --- Registers ---
|
||||
#define MLX5_REG_HOST_ENDIANNESS 0x7004
|
||||
#define MLX5_REG_DTOR 0xC00E
|
||||
|
||||
// --- Misc ---
|
||||
#define MLX5_PCI_CMD_XPORT 0x07
|
||||
#define MLX5_CMD_DATA_BLOCK_SIZE 512
|
||||
#define CMD_OWNER_HW 0x01
|
||||
|
||||
// --- IFC cmd_hca_cap bit offsets ---
|
||||
#define CAP_GEN_ABS_NATIVE_PORT_NUM 0x007
|
||||
#define CAP_GEN_HCA_CAP_2 0x020
|
||||
#define CAP_GEN_EVENT_ON_VHCA_STATE_ALLOCATED 0x023
|
||||
#define CAP_GEN_EVENT_ON_VHCA_STATE_ACTIVE 0x024
|
||||
#define CAP_GEN_EVENT_ON_VHCA_STATE_IN_USE 0x025
|
||||
#define CAP_GEN_EVENT_ON_VHCA_STATE_TEARDOWN_REQUEST 0x026
|
||||
#define CAP_GEN_LOG_MAX_QP 0x09B
|
||||
#define CAP_GEN_LOG_MAX_CQ 0x0DB
|
||||
#define CAP_GEN_RELEASE_ALL_PAGES 0x145
|
||||
#define CAP_GEN_CACHE_LINE_128BYTE 0x164
|
||||
#define CAP_GEN_NUM_PORTS 0x1B8
|
||||
#define CAP_GEN_PKEY_TABLE_SIZE 0x190
|
||||
#define CAP_GEN_PCI_SYNC_FOR_FW_UPDATE_EVENT 0x1F1
|
||||
#define CAP_GEN_CMDIF_CHECKSUM 0x210
|
||||
#define CAP_GEN_DCT 0x21A
|
||||
#define CAP_GEN_ROCE 0x21D
|
||||
#define CAP_GEN_ATOMIC 0x21E
|
||||
#define CAP_GEN_ODP 0x227
|
||||
#define CAP_GEN_MKEY_BY_NAME 0x266
|
||||
#define CAP_GEN_LOG_MAX_PD 0x32B
|
||||
#define CAP_GEN_PCIE_RESET_USING_HOTRESET 0x335
|
||||
#define CAP_GEN_PCI_SYNC_FOR_FW_UPDATE_WITH_DRIVER_UNLOAD 0x336
|
||||
#define CAP_GEN_VHCA_STATE 0x3EA
|
||||
#define CAP_GEN_ROCE_RW_SUPPORTED 0x3A1
|
||||
#define CAP_GEN_LOG_MAX_CURRENT_UC_LIST 0x3FB
|
||||
#define CAP_GEN_LOG_UAR_PAGE_SZ 0x490
|
||||
#define CAP_GEN_NUM_VHCA_PORTS 0x610
|
||||
#define CAP_GEN_SW_OWNER_ID 0x61E
|
||||
#define CAP_GEN_NUM_TOTAL_DYNAMIC_VF_MSIX 0x708
|
||||
Reference in New Issue
Block a user