Restructure iris into host/device split, delete IrisGluon

README.md

@@ -131,7 +131,8 @@ import torch.distributed as dist
 import torch.multiprocessing as mp
 from triton.experimental import gluon
 from triton.experimental.gluon import language as gl
-import iris.experimental.iris_gluon as iris_gl
+import iris
+from iris.gluon import IrisDeviceCtx
 
 # Device-side APIs - context encapsulates heap_bases
 @gluon.jit
@@ -163,20 +164,20 @@ def _worker(rank, world_size):
 
     # Iris initialization
     heap_size = 2**30   # 1GiB symmetric heap
-    iris_ctx = iris_gl.iris(heap_size)
+    iris_ctx = iris.iris(heap_size)
     context_tensor = iris_ctx.get_device_context()  # Get encoded context
     cur_rank = iris_ctx.get_rank()
-    
+
     # Iris tensor allocation
     buffer_size = 4096  # 4K elements buffer
     buffer = iris_ctx.zeros(buffer_size, device="cuda", dtype=torch.float32)
-    
+
     # Launch the kernel on rank 0
     block_size = 1024
     grid = (buffer_size + block_size - 1) // block_size
     source_rank = 0
     if cur_rank == source_rank:
-        kernel[(grid,)](iris_gl.IrisDeviceCtx, context_tensor, 
+        kernel[(grid,)](IrisDeviceCtx, context_tensor,
                        buffer, buffer_size, block_size, num_warps=1)
 
     # Synchronize all ranks

docs/conf.py

@@ -86,15 +86,15 @@
 autodoc_typehints = "description"
 autodoc_typehints_format = "short"
 
-# Render objects without full module path (e.g., show "Iris" instead of "iris.iris.Iris")
+# Render objects without full module path (e.g., show "Iris" instead of "iris.host.iris.Iris")
 add_module_names = False
 
 # Mock heavy/runtime-only dependencies when building docs
 autodoc_mock_imports = [
     "torch",
     "numpy",
-    "iris._distributed_helpers",
-    "iris.hip",
+    "iris.host.distributed.helpers",
+    "iris.host.platform.hip",
     "tritonblas",
 ]
 
@@ -118,6 +118,9 @@ def __call__(self, func=None, **kwargs):
 sys.modules["triton.language"] = triton_language_mock
 sys.modules["triton.language.core"] = MagicMock()
 sys.modules["triton.language.core"]._aggregate = lambda cls: cls  # Preserve class
+sys.modules["triton.language.extra"] = MagicMock()
+sys.modules["triton.language.extra.hip"] = MagicMock()
+sys.modules["triton.language.target_info"] = MagicMock()
 
 
 # Mock triton modules with docstring-preserving jit decorator

docs/index.md

@@ -135,7 +135,8 @@ import torch.distributed as dist
 import torch.multiprocessing as mp
 from triton.experimental import gluon
 from triton.experimental.gluon import language as gl
-import iris.experimental.iris_gluon as iris_gl
+import iris
+from iris.gluon import IrisDeviceCtx
 
 # Device-side APIs - context encapsulates heap_bases
 @gluon.jit
@@ -167,20 +168,20 @@ def _worker(rank, world_size):
 
     # Iris initialization
     heap_size = 2**30   # 1GiB symmetric heap
-    iris_ctx = iris_gl.iris(heap_size)
+    iris_ctx = iris.iris(heap_size)
     context_tensor = iris_ctx.get_device_context()  # Get encoded context
     cur_rank = iris_ctx.get_rank()
-    
+
     # Iris tensor allocation
     buffer_size = 4096  # 4K elements buffer
     buffer = iris_ctx.zeros(buffer_size, device="cuda", dtype=torch.float32)
-    
+
     # Launch the kernel on rank 0
     block_size = 1024
     grid = (buffer_size + block_size - 1) // block_size
     source_rank = 0
     if cur_rank == source_rank:
-        kernel[(grid,)](iris_gl.IrisDeviceCtx, context_tensor, 
+        kernel[(grid,)](IrisDeviceCtx, context_tensor,
                        buffer, buffer_size, block_size, num_warps=1)
 
     # Synchronize all ranks

docs/reference/api-reference.md

@@ -1,25 +1,21 @@
 # API Reference
 
-Explore Iris APIs. The reference is broken down into focused sections to mirror common workflows:
-
-- The `Iris` class itself (constructor and helper utilities)
-- Tensor-like creation methods on the `Iris` context
-- Triton device-side functions for remote memory ops and atomics
-- Collective communication operations (CCL)
-- Fused GEMM + CCL operations
-- Experimental Gluon APIs (using `@aggregate` and `@gluon.jit`)
-
-Use the links below to navigate:
-
-- [Triton](triton/overview.md)
-  - [Iris Class](triton/class.md)
-  - [Tensor Creation](triton/tensor-creation.md)
-  - [Device Functions](triton/device-functions.md)
-  - [Collective Communication (CCL)](triton/ccl.md)
-  - [Fused GEMM + CCL Operations](triton/ops.md)
-- [Gluon (Experimental)](gluon/overview.md)
-  - [Iris Class](gluon/class.md)
-  - [Tensor Creation](gluon/tensor-creation.md)
-  - [Device Functions](gluon/device-functions.md)
-  - [Collective Communication (CCL)](gluon/ccl.md)
+Explore Iris APIs. The host-side API (class, tensor creation, CCL) is identical across all backends. Only the device-side API differs between Triton and Gluon.
+
+## Host API (All Backends)
+
+- [Iris Class](host/class.md)
+- [Tensor Creation](host/tensor-creation.md)
+- [Collective Communication (CCL)](host/ccl.md)
+
+## Triton Backend
+
+- [Overview](triton/overview.md)
+- [Device Functions](triton/device-functions.md)
+- [Fused GEMM + CCL Operations](triton/ops.md)
+
+## Gluon Backend (Experimental)
+
+- [Overview](gluon/overview.md)
+- [Device Functions](gluon/device-functions.md)
 

docs/reference/gluon/device-functions.md

@@ -10,95 +10,95 @@ Device-side functions provided by Iris Gluon for remote memory operations and at
 
 ### initialize
 ```{eval-rst}
-.. automethod:: iris.experimental.iris_gluon.IrisDeviceCtx.initialize
+.. automethod:: iris.device.gluon.context.IrisDeviceCtx.initialize
    :noindex:
 ```
 
 ## Memory transfer operations
 
 ### load
 ```{eval-rst}
-.. automethod:: iris.experimental.iris_gluon.IrisDeviceCtx.load
+.. automethod:: iris.device.gluon.context.IrisDeviceCtx.load
    :noindex:
 ```
 
 ### store
 ```{eval-rst}
-.. automethod:: iris.experimental.iris_gluon.IrisDeviceCtx.store
+.. automethod:: iris.device.gluon.context.IrisDeviceCtx.store
    :noindex:
 ```
 
 ### copy
 ```{eval-rst}
-.. automethod:: iris.experimental.iris_gluon.IrisDeviceCtx.copy
+.. automethod:: iris.device.gluon.context.IrisDeviceCtx.copy
    :noindex:
 ```
 
 ### get
 ```{eval-rst}
-.. automethod:: iris.experimental.iris_gluon.IrisDeviceCtx.get
+.. automethod:: iris.device.gluon.context.IrisDeviceCtx.get
    :noindex:
 ```
 
 ### put
 ```{eval-rst}
-.. automethod:: iris.experimental.iris_gluon.IrisDeviceCtx.put
+.. automethod:: iris.device.gluon.context.IrisDeviceCtx.put
    :noindex:
 ```
 
 ## Atomic operations
 
 ### atomic_add
 ```{eval-rst}
-.. automethod:: iris.experimental.iris_gluon.IrisDeviceCtx.atomic_add
+.. automethod:: iris.device.gluon.context.IrisDeviceCtx.atomic_add
    :noindex:
 ```
 
 ### atomic_sub
 ```{eval-rst}
-.. automethod:: iris.experimental.iris_gluon.IrisDeviceCtx.atomic_sub
+.. automethod:: iris.device.gluon.context.IrisDeviceCtx.atomic_sub
    :noindex:
 ```
 
 ### atomic_cas
 ```{eval-rst}
-.. automethod:: iris.experimental.iris_gluon.IrisDeviceCtx.atomic_cas
+.. automethod:: iris.device.gluon.context.IrisDeviceCtx.atomic_cas
    :noindex:
 ```
 
 ### atomic_xchg
 ```{eval-rst}
-.. automethod:: iris.experimental.iris_gluon.IrisDeviceCtx.atomic_xchg
+.. automethod:: iris.device.gluon.context.IrisDeviceCtx.atomic_xchg
    :noindex:
 ```
 
 ### atomic_xor
 ```{eval-rst}
-.. automethod:: iris.experimental.iris_gluon.IrisDeviceCtx.atomic_xor
+.. automethod:: iris.device.gluon.context.IrisDeviceCtx.atomic_xor
    :noindex:
 ```
 
 ### atomic_and
 ```{eval-rst}
-.. automethod:: iris.experimental.iris_gluon.IrisDeviceCtx.atomic_and
+.. automethod:: iris.device.gluon.context.IrisDeviceCtx.atomic_and
    :noindex:
 ```
 
 ### atomic_or
 ```{eval-rst}
-.. automethod:: iris.experimental.iris_gluon.IrisDeviceCtx.atomic_or
+.. automethod:: iris.device.gluon.context.IrisDeviceCtx.atomic_or
    :noindex:
 ```
 
 ### atomic_min
 ```{eval-rst}
-.. automethod:: iris.experimental.iris_gluon.IrisDeviceCtx.atomic_min
+.. automethod:: iris.device.gluon.context.IrisDeviceCtx.atomic_min
    :noindex:
 ```
 
 ### atomic_max
 ```{eval-rst}
-.. automethod:: iris.experimental.iris_gluon.IrisDeviceCtx.atomic_max
+.. automethod:: iris.device.gluon.context.IrisDeviceCtx.atomic_max
    :noindex:
 ```
 

docs/reference/gluon/overview.md

@@ -26,12 +26,13 @@ The Gluon API provides a Triton Gluon-based implementation of Iris that uses the
 ## Usage Example
 
 ```python
-import iris.experimental.iris_gluon as iris_gl
+import iris
+from iris.gluon import IrisDeviceCtx
 from triton.experimental import gluon
 from triton.experimental.gluon import language as gl
 
 # Host-side: Initialize Iris Gluon context
-ctx = iris_gl.iris(heap_size=2**30)  # 1GB heap
+ctx = iris.iris(heap_size=2**30)  # 1GB heap
 context_tensor = ctx.get_device_context()
 
 # Device-side: Use in Gluon kernels
@@ -49,10 +50,10 @@ def kernel(IrisDeviceCtx: gl.constexpr, context_tensor, buffer):
 
 Explore the API by section:
 
-- [Iris Class](class.md)
-- [Tensor Creation](tensor-creation.md)
+- [Iris Class](../host/class.md)
+- [Tensor Creation](../host/tensor-creation.md)
 - [Device Functions](device-functions.md)
-- [Collective Communication (CCL)](ccl.md)
+- [Collective Communication (CCL)](../host/ccl.md)
 
 ## Complete Example: Producer-Consumer Pattern
 
@@ -64,7 +65,8 @@ import torch.distributed as dist
 import torch.multiprocessing as mp
 from triton.experimental import gluon
 from triton.experimental.gluon import language as gl
-import iris.experimental.iris_gluon as iris_gl
+import iris
+from iris.gluon import IrisDeviceCtx
 
 @gluon.jit
 def producer_kernel(
@@ -137,7 +139,7 @@ def worker(rank, world_size):
     )
 
     # Initialize Iris Gluon
-    ctx = iris_gl.iris(heap_size=2**30)
+    ctx = iris.iris(heap_size=2**30)
     context_tensor = ctx.get_device_context()
 
     # Allocate buffers
@@ -159,7 +161,7 @@ def worker(rank, world_size):
     if rank == producer_rank:
         ctx.info(f"Rank {rank} producing data...")
         producer_kernel[grid](
-            iris_gl.IrisDeviceCtx,
+            IrisDeviceCtx,
             context_tensor,
             source,
             target,
@@ -173,7 +175,7 @@ def worker(rank, world_size):
     else:
         ctx.info(f"Rank {rank} consuming data...")
         consumer_kernel[grid](
-            iris_gl.IrisDeviceCtx,
+            IrisDeviceCtx,
             context_tensor,
             target,
             flag,