Add as_remote(rank) pointer translation API

[mawad-amd]

Summary

Just an idea I'm prototyping — not planning to merge this for now.

as_remote gives a clean way to get a pre-translated pointer to another rank's symmetric heap copy:

• Host-side: ctx.as_remote(tensor, rank) → zero-copy torch.Tensor whose data_ptr() points to the target rank's copy. Same shape/dtype/strides.
• Device-side (Gluon): ctx.as_remote(ptr, rank) → translated pointer for direct use with gl.load/gl.store.

Useful for hoisting translation out of loops or passing pre-translated pointers to kernels without manual heap base arithmetic.

Host-side example

ctx = iris.iris(heap_size=2**30)
buf = ctx.zeros(1024, dtype=torch.float32)

# Get a view pointing to rank 3's copy of buf
remote_buf = ctx.as_remote(buf, rank=3)

# remote_buf has same shape/dtype/strides, but data_ptr()
# points into rank 3's symmetric heap
print(remote_buf.shape)       # torch.Size([1024])
print(remote_buf.data_ptr())  # address in rank 3's heap

Device-side example (Gluon)

@gluon.jit
def read_from_peer(IrisDeviceCtx: gl.constexpr, context_tensor, buf, out, peer: gl.constexpr, BLOCK: gl.constexpr):
    ctx = IrisDeviceCtx.initialize(context_tensor)
    offs = gl.program_id(0) * BLOCK + gl.arange(0, BLOCK)
    mask = offs < BLOCK

    # Translate pointer once, then use gl.load/gl.store directly
    remote_ptr = ctx.as_remote(buf + offs, peer)
    data = gl.load(remote_ptr, mask=mask)
    gl.store(out + offs, data, mask=mask)

Test plan

• Host-side: basic, pointer math, self-rank, validation errors, non-contiguous, multi-dtype
• Device-side: read via as_remote + gl.load, write via as_remote + gl.store
• Run on real GPU hardware with torchrun --nproc_per_node=4

🤖 Generated with Claude Code