Improve docs

[mawad-amd]

First attempt at improving docs. local/remote is still a bit confusing.

Closes #46

• Translate should be private
• Improve docs

[Copilot]

Pull Request Overview

This PR refactors the internal pointer translation helper, tightens up public exports, and enriches docstrings for remote memory operations.

• Renames translate to private __translate and updates all function signatures to use local_ptr/local_rank/remote_rank conventions
• Enhances docstrings for load, store, get, put, and atomic operations with detailed parameter and return descriptions
• Removes translate from public exports in __init__.py

Reviewed Changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
iris/iris.py	Renamed translate to __translate, updated signatures and improved docstrings
iris/init.py	Removed translate from imports and __all__ to make it an internal helper

Comments suppressed due to low confidence (3)

iris/iris.py:318

• Add a docstring for __translate explaining its purpose, parameters, and return value to maintain consistency with other Triton helper functions.

def __translate(local_ptr, local_rank, remote_rank, heap_bases, debug=False):

iris/init.py:11

• Removing translate from the public exports is a breaking change; consider deprecating it for one release or updating the changelog to notify users of this API change.

    iris,

iris/iris.py:318

• The new pointer translation logic in __translate (and its use in load/store/get/put) should have unit tests to verify correctness across different rank configurations.

def __translate(local_ptr, local_rank, remote_rank, heap_bases, debug=False):

[neoblizz]

Why do we have to describe it as local/remote only? If we use iris.load it should do local-load if your source and destination ranks are the same, correct?

[mawad-amd]

Correct. It is confusing and I would like to resolve that. Do you have suggestions?

I have a few:

# 1. Emphasizes direction of data flow
def load(pointer, to_rank, from_rank, heap_bases, mask=None):

# 2. Uses 'cur' to indicate the calling rank
def load(pointer, cur_rank, dst_rank, heap_bases, mask=None):

# 3. Generalizes roles as caller/target
def load(pointer, caller_rank, target_rank, heap_bases, mask=None):

The pointer argument can be named local_ptr, sym_ptr, address, or pointer.

I do like the from/to at the moment. Here are all APIs:

def load(pointer, to_rank, from_rank, heap_bases, mask=None):

def store(pointer, val, from_rank, to_rank, heap_bases, mask=None):

def get(from_pointer, to_pointer, to_rank, from_rank, heap_bases, mask=None):

def put(from_pointer, to_pointer, from_rank, to_rank, heap_bases, mask=None):

def atomic_add(pointer, val, from_rank, to_rank, heap_bases, mask=None, sem=None, scope=None):

def atomic_sub(pointer, val, from_rank, to_rank, heap_bases, mask=None, sem=None, scope=None):

def atomic_cas(pointer, cmp, val, from_rank, to_rank, heap_bases, sem=None, scope=None):

def atomic_xchg(pointer, val, from_rank, to_rank, heap_bases, mask=None, sem=None, scope=None):

[neoblizz]

Correct. It is confusing and I would like to resolve that. Do you have suggestions?

I have a few:

# 1. Emphasizes direction of data flow
def load(pointer, to_rank, from_rank, heap_bases, mask=None):

# 2. Uses 'cur' to indicate the calling rank
def load(pointer, cur_rank, dst_rank, heap_bases, mask=None):

# 3. Generalizes roles as caller/target
def load(pointer, caller_rank, target_rank, heap_bases, mask=None):

The pointer argument can be named local_ptr, sym_ptr, address, or pointer.

I do like the from/to at the moment. Here are all APIs:

def load(pointer, to_rank, from_rank, heap_bases, mask=None):

def store(pointer, val, from_rank, to_rank, heap_bases, mask=None):

def get(from_pointer, to_pointer, to_rank, from_rank, heap_bases, mask=None):

def put(from_pointer, to_pointer, from_rank, to_rank, heap_bases, mask=None):

def atomic_add(pointer, val, from_rank, to_rank, heap_bases, mask=None, sem=None, scope=None):

def atomic_sub(pointer, val, from_rank, to_rank, heap_bases, mask=None, sem=None, scope=None):

def atomic_cas(pointer, cmp, val, from_rank, to_rank, heap_bases, sem=None, scope=None):

def atomic_xchg(pointer, val, from_rank, to_rank, heap_bases, mask=None, sem=None, scope=None):

Notes

• I do like to/from, but not for the from_pointer field.
• Is there a notion of "triggering" a load/store/put/get from a rank thats not "current" or "remote" rank? If so, current cannot be used.

Other suggestions:

• Use value instead of val,
• Use semantics instead of sem
• The problematic (imo) APIs are:

def get(from_pointer, to_pointer, to_rank, from_rank, heap_bases, mask=None):
def put(from_pointer, to_pointer, from_rank, to_rank, heap_bases, mask=None):

[mawad-amd]

Is there a notion of "triggering" a load/store/put/get from a rank thats not "current" or "remote" rank? If so, current cannot be used.

No. For that reason, I was considering adding a text to warn against that e.g., for load "to_rank must be the same rank issuing the operation." But if there is some an interesting use-case we can accommodate that.
Eventually the "current rank" should be implicit and removed, alongside the heap_bases.

The shortened val and sem are to match Triton but I am okay with either (and I slightly prefer spelling out the complete word).

How about this:

def load(pointer, to_rank, from_rank, heap_bases, mask=None):
def store(pointer, val, from_rank, to_rank, heap_bases, mask=None):
def atomic_add(pointer, val, from_rank, to_rank, heap_bases, mask=None, sem=None, scope=None):
def atomic_sub(pointer, val, from_rank, to_rank, heap_bases, mask=None, sem=None, scope=None):
def atomic_cas(pointer, cmp, val, from_rank, to_rank, heap_bases, sem=None, scope=None):
def atomic_xchg(pointer, val, from_rank, to_rank, heap_bases, mask=None, sem=None, scope=None):

and,

def get(dst_ptr, src_ptr, to_rank, from_rank, heap_bases, mask=None):
def put(dst_ptr, src_ptr, to_rank, from_rank, heap_bases, mask=None):

[neoblizz]

No. For that reason, I was considering adding a text to warn against that e.g., for load "to_rank must be the same rank issuing the operation." But if there is some an interesting use-case we can accommodate that.
Eventually the "current rank" should be implicit and removed, alongside the heap_bases.

I am not sure I understand the warning. You mean from_rank must be the same rank issuing the op?
Separately, I think it can be true where thats not the case. Where we can have one GPU thread somehow initiate a copy/move. Lets discuss in a call separately with @BKP.

The shortened val and sem are to match Triton but I am okay with either (and I slightly prefer spelling out the complete word).

How about this:

def load(pointer, to_rank, from_rank, heap_bases, mask=None):
def store(pointer, val, from_rank, to_rank, heap_bases, mask=None):
def atomic_add(pointer, val, from_rank, to_rank, heap_bases, mask=None, sem=None, scope=None):
def atomic_sub(pointer, val, from_rank, to_rank, heap_bases, mask=None, sem=None, scope=None):
def atomic_cas(pointer, cmp, val, from_rank, to_rank, heap_bases, sem=None, scope=None):
def atomic_xchg(pointer, val, from_rank, to_rank, heap_bases, mask=None, sem=None, scope=None):

and,

def get(dst_ptr, src_ptr, to_rank, from_rank, heap_bases, mask=None):
def put(dst_ptr, src_ptr, to_rank, from_rank, heap_bases, mask=None):

I prefer spelled out as well --- triton is not consistent with their naming, you can easily find examples of these.

I like this! (with sem, val, maybe even cmp fully spelled out. I know its minor, but sem could be semaphore or semantics or something else, idk)

[neoblizz]

Still lots of remote, lets discuss the best way to word some of the descriptions.

[neoblizz]

Suggested change

	Loads a value from a remote rank's memory location.
	Loads a value stored from a pointer of the specified rank.

[neoblizz]

Suggested change

	data from the remote memory location.
	data from the remote memory location. If `to_rank` is the same as `from_rank`,
	this function performs a local load operation instead.

[neoblizz]

Suggested change

	to_rank (int): The rank ID to which the pointer will be translated. Must be the current rank where the pointer is local.
	to_rank (int): The rank ID for the pointer where the load will occur. `to_rank` must be the rank where the pointer resides.

[neoblizz]

Suggested change

	from_rank (int): The rank ID from which the pointer originates. Must be the current rank where the pointer is local.
	from_rank (int): The rank ID from which the pointer originates. `from_rank` must be the rank where the pointer resides.

[neoblizz]

Use of the word remote again.