implement JIT lowering for `comm.nccl` ops - #3122

Draft

mofeing wants to merge 14 commits into

mainfrom

ss/comm-nccl-lower-to-jit

Commits on Sep 18, 2026

implement jit lowering for `NcclCommCountOp`
mofeing
committed
9506062
implement jit lowering for `NcclCommUserRankOp`
mofeing
committed
f9075e9
implement jit lowering for `NcclSendOp`
mofeing
committed
f9f9ede
fix datatype on `NcclSendOp` jit lowering
mofeing
committed
c8cff22
fix llvm function type
mofeing
committed
815f656
implement jit lowering of `NcclRecvOp`
mofeing
committed
5d4b4d2
fix error throwing
mofeing
committed
7e4014c
implement jit lowering for `NcclBroadcastOp`
mofeing
committed
d03b5a4
implement jit lowering for `NcclAllReduceOp`
mofeing
committed
e75eb42
several fixes
mofeing
committed
d9cfa21
Add back in enzymexla.device_abi = "cuda" wrapper func attr

romanlee
authored and
mofeing
committed
d6da40e
arg0 isn't a pointer to device data but rather to an slot which contains

romanlee
authored and
mofeing
committed
de1c3e0
Actually, I believe fix in previous commit applies to all jitwrap args

romanlee
authored and
mofeing
committed
9b41d03
conditionally compile code only if cuda is enabled
mofeing
committed
3029b0c