Platform.Linux.Nvidia.PlanHostPeers

1module Platform.Linux.Nvidia.PlanHostPeers
2
3import Data.Bytes
4import Platform.Linux.Nvidia.PlanHost
5import Runtime.NativePhysicalProgram
6import Std.Natural
7
8-- The cards of one process reading and writing one another's memory
9-- (docs/PRD-8XH100-ON-A-BUDGET.md F4), on the direct-RM lifecycle and
10-- PlanHost's sessions (one RM client, a session a card), without UVM: some
11-- RunPod hosts refuse /dev/nvidia-uvm (docs/PLAN-NANOCHAT-SPEEDRUN.md).
12--
13--   1. For every pair of sessions a < b on different cards (`card s` is
14--      session s's), an NV50_P2P object (class 0x503b,
15--      its parent the client: resource_list.h) joining their subdevices,
16--      the peer IDs RM's to choose (cl503b.h).  With it RM programs the
17--      link -- NVLink through the NVSwitch on an HGX H100 -- and gives
18--      each GPU a peer ID for the other (p2p_api.c).
19--   2. For every session s and every session t, t = s included: a virtual
20--      reservation in s's address space at `window t` (NV01_MEMORY_VIRTUAL,
21--      0x70, as PlanHost's own buffers take) and t's buffer `index` mapped
22--      into it (NV_ESC_RM_MAP_MEMORY_DMA, NVOS46: hDevice s's, hMemory t's).
23--      A peer's map carries NVOS46_FLAGS_P2P_ENABLE_NOSLI with source
24--      subdevice 0 (nvos.h; dmaAllocMapping takes the peer ID from the
25--      memory's GPU: virt_mem_allocator_gm107.c); a session's own, or a
26--      session's on the same card (virtual ranks: one GPU, a device object
27--      a session, which RM refuses only on a vGPU: device.c), the plain map.
28--
29-- So every card sees every card's buffer at the same address, `window t`:
30-- a kernel addresses rank t's copy of an element at window t plus the
31-- element's offset, on every rank alike (Realization.Nvidia.SM86.
32-- PeerCollectiveSM86).  The sources were read at open-gpu-kernel-modules
33-- 580.95.05 (kept under /mnt/donto-data/donto-resources).
34
35-- NVOS46_FLAGS_P2P_ENABLE (21:20) = NOSLI (2); P2P_SUBDEVICE_ID (24:22) 0
36def ppPeerFlag : Nat = 0x200000
37-- a pair's P2P object, a reservation's handle
38def ppPairHandle = (lambda unrestricted a : Nat . (lambda unrestricted b : Nat .
39  (naturalAdd 0xB2000000 (naturalAdd (naturalMultiply a 16) b))))
40def ppWindowHandle = (lambda unrestricted session : Nat . (lambda unrestricted peer : Nat .
41  (phSessionHandle session (naturalAdd 0xB3000000 peer))))
42-- NV503B_ALLOC_PARAMETERS: hSubDevice (0), hPeerSubDevice (4), the masks,
43-- mailbox and flags zero; 72 bytes through 580, 88 from 590 (two fabric DMA
44-- infos added: cl503b.h at 590.44.01 and 595.45.04)
45def ppPairParamsSize = (lambda unrestricted driver : (family NvidiaDriverInterface) .
46  (naturalSelect (naturalLess (nvidiaDriverInterfaceMajorOf driver) 590) 72 88))
47-- the call blocks: the UVM DMA selector's (unused under direct RM), the
48-- parameters after its own, the statuses in RM slots 20 and 21
49def ppBlock : Nat = phDMAObject
50def ppParams : Nat = phDMAObjectP
51def ppAllocStatus : Nat = (phRMStatusSlot 20)
52def ppMapStatus : Nat = (phRMStatusSlot 21)
53
54def ppPair = (lambda unrestricted abi : (family NvidiaPlanHostABI) . (lambda unrestricted a : Nat . (lambda unrestricted b : Nat .
55  (phDriverChoice (phDrivers abi)
56    (lambda unrestricted driver : (family NvidiaDriverInterface) . (lambda unrestricted tail : (family NativePhysicalCommands) .
57      (phFill ppParams
58        (phSet32 (phSet32 (phZeros (ppPairParamsSize driver)) 0 (phHSubdeviceS a)) 4 (phHSubdeviceS b))
59        (phAlloc b"alloc:p2p" 0 ppBlock 0 0x503b ppParams (ppPairParamsSize driver) (ppPairHandle a b) ppAllocStatus
60          tail))))))))
61
62-- a choice of two byte strings (flag nonzero: the first)
63def ppSelectBytes = (lambda unrestricted flag : Nat . (lambda unrestricted whenNonzero : Bytes . (lambda unrestricted whenZero : Bytes .
64  (nat-eliminate (lambda unrestricted current : Nat . Bytes) whenZero
65    (lambda unrestricted p : Nat . (lambda unrestricted ignored : Bytes . whenNonzero)) flag))))
66
67-- NVOS46 into session s's reservation of `peer`'s window, as PlanHost's
68-- phMapDMA builds it (hClient copied from the root at run time)
69def ppMapWith =
70  (lambda unrestricted abi : (family NvidiaPlanHostABI) .
71    (lambda unrestricted sameCard : Nat .
72    (lambda unrestricted session : Nat .
73      (lambda unrestricted peer : Nat .
74        (lambda unrestricted index : Nat .
75          (lambda unrestricted gpu : Nat .
76            (lambda unrestricted extent : Nat .
77              (lambda unrestricted driver : (family NvidiaDriverInterface) .
78                (lambda unrestricted tail : (family NativePhysicalCommands) .
79                  (let unrestricted status = (naturalAdd ppBlock (nvidiaDriverInterfaceNVOS46StatusOf driver)) in
80                  (let unrestricted flags = (naturalAdd (phDirectDMAFlags abi) (naturalSelect sameCard 0 ppPeerFlag)) in
81                  (let unrestricted handles =
82                    (phSet32 (phSet32 (phSet32 (phSet32 (phZeros (nvidiaDriverInterfaceNVOS46SizeOf driver))
83                      4 (phHDeviceS session)) 8 (ppWindowHandle session peer)) 12 (phHMemoryS peer index)) 32 flags) in
84                  (let unrestricted block =
85                    (phSet32 (phSet64 (phSet64 handles 24 extent) (nvidiaDriverInterfaceNVOS46DMAOffsetOf driver) gpu)
86                      (nvidiaDriverInterfaceNVOS46StatusOf driver) phStatusSentinel) in
87                  (phCopy ppBlock block
88                    (phCopyStateToState ppBlock (naturalAdd phRoot 8) 4
89                      (phCall phSysIoctl
90                        (phArgs3 (phSlot phSlotControl) (phImm (nvidiaDriverInterfaceMapDMAIoctlOf driver)) (phState ppBlock))
91                        b"" phDiscard
92                        (phCopyStateToState ppMapStatus status 4
93                          (phCheckStateZero32 (ppSelectBytes sameCard b"dma-map-window" b"dma-map-peer")
94                            0 status tail))))))))))))))))))
95def ppMap =
96  (lambda unrestricted abi : (family NvidiaPlanHostABI) .
97    (lambda unrestricted sameCard : Nat .
98    (lambda unrestricted session : Nat .
99      (lambda unrestricted peer : Nat .
100        (lambda unrestricted index : Nat .
101          (lambda unrestricted gpu : Nat .
102            (lambda unrestricted extent : Nat .
103              (phDriverChoice (phDrivers abi) (ppMapWith abi sameCard session peer index gpu extent)))))))))
104
105-- session s's window of `peer`: the reservation, then the map
106def ppWindow =
107  (lambda unrestricted abi : (family NvidiaPlanHostABI) .
108    (lambda unrestricted sameCard : Nat .
109    (lambda unrestricted session : Nat .
110      (lambda unrestricted peer : Nat .
111        (lambda unrestricted index : Nat .
112          (lambda unrestricted gpu : Nat .
113            (lambda unrestricted extent : Nat .
114              (lambda unrestricted tail : (family NativePhysicalCommands) .
115                (phFill ppParams
116                  (phSet32 (phSet64 (phSet64 (phZeros 24) 0 gpu) 8 (naturalSaturatingSubtract (naturalAdd gpu extent) 1)) 16 (phHVASpaceS session))
117                  (phAlloc b"alloc:window" 0 ppBlock (phHDeviceS session) 0x70 ppParams 24 (ppWindowHandle session peer) ppAllocStatus
118                    (ppMap abi sameCard session peer index gpu extent tail)))))))))))
119
120-- Every pair joined, then every session's window of every session:
121-- `sessions` cards, buffer `index` of each (the same index in every
122-- session's layout) of `extent` bytes, peer t's window at `window t`.
123def ppPeers = (lambda unrestricted abi : (family NvidiaPlanHostABI) .
124  (lambda unrestricted card : (pi unrestricted session : Nat . Nat) .
125    (lambda unrestricted sessions : Nat . (lambda unrestricted index : Nat . (lambda unrestricted extent : Nat .
126          (lambda unrestricted window : (pi unrestricted peer : Nat . Nat) .
127            (lambda unrestricted tail : (family NativePhysicalCommands) .
128              (let unrestricted pairs =
129                (nat-eliminate (lambda unrestricted n : Nat . (pi unrestricted after : (family NativePhysicalCommands) . (family NativePhysicalCommands)))
130                  (lambda unrestricted after : (family NativePhysicalCommands) . after)
131                  (lambda unrestricted b : Nat . (lambda unrestricted earlier : (pi unrestricted after : (family NativePhysicalCommands) . (family NativePhysicalCommands)) .
132                      (lambda unrestricted after : (family NativePhysicalCommands) .
133                        (earlier
134                          (nat-eliminate (lambda unrestricted n : Nat . (family NativePhysicalCommands)) after
135                            (lambda unrestricted a : Nat . (lambda unrestricted later : (family NativePhysicalCommands) .
136                                (nat-eliminate (lambda unrestricted n : Nat . (family NativePhysicalCommands)) (ppPair abi a b later)
137                                  (lambda unrestricted q : Nat . (lambda unrestricted ignored : (family NativePhysicalCommands) . later))
138                                  (naturalEqual (card a) (card b)))))
139                            b)))))
140                  sessions) in
141                (let unrestricted windows =
142                  (nat-eliminate (lambda unrestricted n : Nat . (family NativePhysicalCommands)) tail
143                    (lambda unrestricted k : Nat . (lambda unrestricted later : (family NativePhysicalCommands) .
144                        (let unrestricted session = (naturalDivideUnchecked k sessions) in
145                          (let unrestricted peer = (naturalModuloUnchecked k sessions) in
146                            (ppWindow abi (naturalEqual (card session) (card peer)) session peer index (window peer) extent later)))))
147                    (naturalMultiply sessions sessions)) in
148                  (pairs windows))))))))))