Add iris.concurrent (concurrent, independent, fused comm+gemm)

.gitignore

@@ -69,3 +69,6 @@ docker/rocm-systems/
 .intellikit
 .github/agents/skills/
 docs/benchmark-results/*.png
+
+# Committed concurrent tuning corpus (durable data, see iris/concurrent/tuning_data/README.md)
+!iris/concurrent/tuning_data/candidate_set.json

README.md

@@ -59,6 +59,7 @@ import triton
 import triton.language as tl
 import iris
 
+
 # Device-side APIs
 @triton.jit
 def kernel(buffer, buffer_size: tl.constexpr, block_size: tl.constexpr, heap_bases_ptr):
@@ -73,9 +74,8 @@ def kernel(buffer, buffer_size: tl.constexpr, block_size: tl.constexpr, heap_bas
     # Store 1 in the target buffer at each offset
     source_rank = 0
     target_rank = 1
-    iris.store(buffer + offsets, 1,
-            source_rank, target_rank,
-            heap_bases_ptr, mask=mask)
+    iris.store(buffer + offsets, 1, source_rank, target_rank, heap_bases_ptr, mask=mask)
+
 
 def _worker(rank, world_size):
     # Torch distributed initialization
@@ -85,11 +85,11 @@ def _worker(rank, world_size):
         rank=rank,
         world_size=world_size,
         init_method="tcp://127.0.0.1:29500",
-        device_id=torch.device(f"cuda:{device_id}")
+        device_id=torch.device(f"cuda:{device_id}"),
     )
 
     # Iris initialization
-    heap_size = 2**30   # 1GiB symmetric heap for inter-GPU communication
+    heap_size = 2**30  # 1GiB symmetric heap for inter-GPU communication
     iris_ctx = iris.iris(heap_size)
     cur_rank = iris_ctx.get_rank()
 
@@ -113,6 +113,7 @@ def _worker(rank, world_size):
     iris_ctx.barrier()
     dist.destroy_process_group()
 
+
 if __name__ == "__main__":
     world_size = 2  # Using two ranks
     mp.spawn(_worker, args=(world_size,), nprocs=world_size, join=True)
@@ -134,13 +135,13 @@ from triton.experimental.gluon import language as gl
 import iris
 from iris.gluon import IrisDeviceCtx
 
+
 # Device-side APIs - context encapsulates heap_bases
 @gluon.jit
-def kernel(IrisDeviceCtx: gl.constexpr, context_tensor,
-          buffer, buffer_size: gl.constexpr, block_size: gl.constexpr):
+def kernel(IrisDeviceCtx: gl.constexpr, context_tensor, buffer, buffer_size: gl.constexpr, block_size: gl.constexpr):
     # Initialize device context from tensor
     ctx = IrisDeviceCtx.initialize(context_tensor)
-    
+
     pid = gl.program_id(0)
     block_start = pid * block_size
     layout: gl.constexpr = gl.BlockedLayout([1], [64], [1], [0])
@@ -151,6 +152,7 @@ def kernel(IrisDeviceCtx: gl.constexpr, context_tensor,
     target_rank = 1
     ctx.store(buffer + offsets, 1, target_rank, mask=mask)
 
+
 def _worker(rank, world_size):
     # Torch distributed initialization
     device_id = rank % torch.cuda.device_count()
@@ -159,11 +161,11 @@ def _worker(rank, world_size):
         rank=rank,
         world_size=world_size,
         init_method="tcp://127.0.0.1:29500",
-        device_id=torch.device(f"cuda:{device_id}")
+        device_id=torch.device(f"cuda:{device_id}"),
     )
 
     # Iris initialization
-    heap_size = 2**30   # 1GiB symmetric heap
+    heap_size = 2**30  # 1GiB symmetric heap
     iris_ctx = iris.iris(heap_size)
     context_tensor = iris_ctx.get_device_context()  # Get encoded context
     cur_rank = iris_ctx.get_rank()
@@ -177,13 +179,13 @@ def _worker(rank, world_size):
     grid = (buffer_size + block_size - 1) // block_size
     source_rank = 0
     if cur_rank == source_rank:
-        kernel[(grid,)](IrisDeviceCtx, context_tensor,
-                       buffer, buffer_size, block_size, num_warps=1)
+        kernel[(grid,)](IrisDeviceCtx, context_tensor, buffer, buffer_size, block_size, num_warps=1)
 
     # Synchronize all ranks
     iris_ctx.barrier()
     dist.destroy_process_group()
 
+
 if __name__ == "__main__":
     world_size = 2  # Using two ranks
     mp.spawn(_worker, args=(world_size,), nprocs=world_size, join=True)

docs/conceptual/programming-model.md

@@ -149,6 +149,7 @@ def load(pointer, to_rank, from_rank, heap_bases, mask=None):
     result = tl.load(translated_ptr, mask=mask)
     return result
 
+
 @triton.jit
 def __translate(ptr, from_rank, to_rank, heap_bases):
     from_base = tl.load(heap_bases + from_rank)

docs/conceptual/taxonomy.md

@@ -120,15 +120,13 @@ The fundamental insight is that **achieving peak performance requires computatio
 # Producer side
 compute_tile(...)
 # Release pattern - notify tile is ready
-iris.atomic_cas(flag + tile_id, 0, 1, consumer_rank, 
-                sem="release", scope="sys")
+iris.atomic_cas(flag + tile_id, 0, 1, consumer_rank, sem="release", scope="sys")
 
 # Consumer side
 # Acquire pattern - wait for tile
 done = 0
 while done == 0:
-    done = iris.atomic_cas(flag + tile_id, 1, 0, consumer_rank,
-                          sem="acquire", scope="sys")
+    done = iris.atomic_cas(flag + tile_id, 1, 0, consumer_rank, sem="acquire", scope="sys")
 # Consume data
 data = iris.load(buffer + offsets, consumer_rank, mask=mask)
 ```

docs/index.md

@@ -66,6 +66,7 @@ import triton
 import triton.language as tl
 import iris
 
+
 # Device-side APIs
 @triton.jit
 def kernel(buffer, buffer_size: tl.constexpr, block_size: tl.constexpr, heap_bases_ptr):
@@ -80,9 +81,8 @@ def kernel(buffer, buffer_size: tl.constexpr, block_size: tl.constexpr, heap_bas
     # Store 1 in the target buffer at each offset
     source_rank = 0
     target_rank = 1
-    iris.store(buffer + offsets, 1,
-            source_rank, target_rank,
-            heap_bases_ptr, mask=mask)
+    iris.store(buffer + offsets, 1, source_rank, target_rank, heap_bases_ptr, mask=mask)
+
 
 def _worker(rank, world_size):
     # Torch distributed initialization
@@ -92,11 +92,11 @@ def _worker(rank, world_size):
         rank=rank,
         world_size=world_size,
         init_method="tcp://127.0.0.1:29500",
-        device_id=torch.device(f"cuda:{device_id}")
+        device_id=torch.device(f"cuda:{device_id}"),
     )
 
     # Iris initialization
-    heap_size = 2**30   # 1GiB symmetric heap for inter-GPU communication
+    heap_size = 2**30  # 1GiB symmetric heap for inter-GPU communication
     iris_ctx = iris.iris(heap_size)
     cur_rank = iris_ctx.get_rank()
 
@@ -120,6 +120,7 @@ def _worker(rank, world_size):
     iris_ctx.barrier()
     dist.destroy_process_group()
 
+
 if __name__ == "__main__":
     world_size = 2  # Using two ranks
     mp.spawn(_worker, args=(world_size,), nprocs=world_size, join=True)
@@ -138,13 +139,13 @@ from triton.experimental.gluon import language as gl
 import iris
 from iris.gluon import IrisDeviceCtx
 
+
 # Device-side APIs - context encapsulates heap_bases
 @gluon.jit
-def kernel(IrisDeviceCtx: gl.constexpr, context_tensor,
-          buffer, buffer_size: gl.constexpr, block_size: gl.constexpr):
+def kernel(IrisDeviceCtx: gl.constexpr, context_tensor, buffer, buffer_size: gl.constexpr, block_size: gl.constexpr):
     # Initialize device context from tensor
     ctx = IrisDeviceCtx.initialize(context_tensor)
-    
+
     pid = gl.program_id(0)
     block_start = pid * block_size
     layout: gl.constexpr = gl.BlockedLayout([1], [64], [1], [0])
@@ -155,6 +156,7 @@ def kernel(IrisDeviceCtx: gl.constexpr, context_tensor,
     target_rank = 1
     ctx.store(buffer + offsets, 1, target_rank, mask=mask)
 
+
 def _worker(rank, world_size):
     # Torch distributed initialization
     device_id = rank % torch.cuda.device_count()
@@ -163,11 +165,11 @@ def _worker(rank, world_size):
         rank=rank,
         world_size=world_size,
         init_method="tcp://127.0.0.1:29500",
-        device_id=torch.device(f"cuda:{device_id}")
+        device_id=torch.device(f"cuda:{device_id}"),
     )
 
     # Iris initialization
-    heap_size = 2**30   # 1GiB symmetric heap
+    heap_size = 2**30  # 1GiB symmetric heap
     iris_ctx = iris.iris(heap_size)
     context_tensor = iris_ctx.get_device_context()  # Get encoded context
     cur_rank = iris_ctx.get_rank()
@@ -181,13 +183,13 @@ def _worker(rank, world_size):
     grid = (buffer_size + block_size - 1) // block_size
     source_rank = 0
     if cur_rank == source_rank:
-        kernel[(grid,)](IrisDeviceCtx, context_tensor,
-                       buffer, buffer_size, block_size, num_warps=1)
+        kernel[(grid,)](IrisDeviceCtx, context_tensor, buffer, buffer_size, block_size, num_warps=1)
 
     # Synchronize all ranks
     iris_ctx.barrier()
     dist.destroy_process_group()
 
+
 if __name__ == "__main__":
     world_size = 2  # Using two ranks
     mp.spawn(_worker, args=(world_size,), nprocs=world_size, join=True)

docs/reference/gluon/overview.md

@@ -35,12 +35,13 @@ from triton.experimental.gluon import language as gl
 ctx = iris.iris(heap_size=2**30)  # 1GB heap
 context_tensor = ctx.get_device_context()
 
+
 # Device-side: Use in Gluon kernels
 @gluon.jit
 def kernel(IrisDeviceCtx: gl.constexpr, context_tensor, buffer):
     # Initialize device context from tensor
     ctx = IrisDeviceCtx.initialize(context_tensor)
-    
+
     # Perform remote memory operations
     data = ctx.load(buffer, from_rank=1)
     ctx.store(buffer, data, to_rank=0)
@@ -68,6 +69,7 @@ from triton.experimental.gluon import language as gl
 import iris
 from iris.gluon import IrisDeviceCtx
 
+
 @gluon.jit
 def producer_kernel(
     IrisDeviceCtx: gl.constexpr,
@@ -82,21 +84,22 @@ def producer_kernel(
 ):
     ctx = IrisDeviceCtx.initialize(context_tensor)
     pid = gl.program_id(0)
-    
+
     block_start = pid * BLOCK_SIZE
     layout: gl.constexpr = gl.BlockedLayout([1], [64], [1], [0])
     offsets = block_start + gl.arange(0, BLOCK_SIZE, layout=layout)
     mask = offsets < buffer_size
-    
+
     # Load from producer's buffer
     values = ctx.load(source_buffer + offsets, producer_rank, mask=mask)
-    
+
     # Store to consumer's buffer
     ctx.store(target_buffer + offsets, values, consumer_rank, mask=mask)
-    
+
     # Signal completion
     ctx.atomic_cas(flag + pid, 0, 1, consumer_rank, sem="release", scope="sys")
 
+
 @gluon.jit
 def consumer_kernel(
     IrisDeviceCtx: gl.constexpr,
@@ -109,53 +112,56 @@ def consumer_kernel(
 ):
     ctx = IrisDeviceCtx.initialize(context_tensor)
     pid = gl.program_id(0)
-    
+
     block_start = pid * BLOCK_SIZE
     layout: gl.constexpr = gl.BlockedLayout([1], [64], [1], [0])
     offsets = block_start + gl.arange(0, BLOCK_SIZE, layout=layout)
     mask = offsets < buffer_size
-    
+
     # Wait for producer
     done = 0
     while done == 0:
         done = ctx.atomic_cas(flag + pid, 1, 0, consumer_rank, sem="acquire", scope="sys")
-    
+
     # Read from buffer
     values = ctx.load(buffer + offsets, consumer_rank, mask=mask)
-    
+
     # Process values...
     values = values * 2
-    
+
     # Store back
     ctx.store(buffer + offsets, values, consumer_rank, mask=mask)
 
+
 def worker(rank, world_size):
     # Initialize distributed
     device_id = rank % torch.cuda.device_count()
     dist.init_process_group(
-        backend="nccl", rank=rank, world_size=world_size,
+        backend="nccl",
+        rank=rank,
+        world_size=world_size,
         init_method="tcp://127.0.0.1:29500",
-        device_id=torch.device(f"cuda:{device_id}")
+        device_id=torch.device(f"cuda:{device_id}"),
     )
-    
+
     # Initialize Iris Gluon
     ctx = iris.iris(heap_size=2**30)
     context_tensor = ctx.get_device_context()
-    
+
     # Allocate buffers
     buffer_size = 1024
     block_size = 256
     source = ctx.zeros(buffer_size, dtype=torch.float32)
     target = ctx.zeros(buffer_size, dtype=torch.float32)
     num_blocks = triton.cdiv(buffer_size, block_size)
     flag = ctx.zeros(num_blocks, dtype=torch.int32)
-    
+
     # Initialize source data on producer
     producer_rank = 0
     consumer_rank = 1
     if rank == producer_rank:
         source.fill_(42.0)
-    
+
     # Launch kernels based on rank
     grid = (num_blocks,)
     if rank == producer_rank:
@@ -184,20 +190,21 @@ def worker(rank, world_size):
             block_size,
             num_warps=1,
         )
-    
+
     ctx.barrier()
-    
+
     # Validate on consumer
     if rank == consumer_rank:
         expected = source * 2  # Consumer doubles the values
         if torch.allclose(target, expected, atol=1):
             ctx.info("Validation successful!")
         else:
             ctx.error("Validation failed!")
-    
+
     ctx.barrier()
     dist.destroy_process_group()
 
+
 if __name__ == "__main__":
     world_size = 2
     mp.spawn(worker, args=(world_size,), nprocs=world_size, join=True)

docs/reference/triton/overview.md

@@ -16,13 +16,14 @@ import triton.language as tl
 # Host-side: Initialize Iris context
 ctx = iris.iris(heap_size=2**30)  # 1GB heap
 
+
 # Device-side: Use in Triton kernels
 @triton.jit
 def kernel(ptr, heap_bases, cur_rank, remote_rank, BLOCK_SIZE: tl.constexpr):
     pid = tl.program_id(0)
     block_start = pid * BLOCK_SIZE
     offsets = block_start + tl.arange(0, BLOCK_SIZE)
-    
+
     # Perform remote memory operations
     data = iris.load(ptr + offsets, cur_rank, remote_rank, heap_bases)
     iris.store(ptr + offsets, data, cur_rank, remote_rank, heap_bases)