[TLE] Move TLE distributed API under language namespace (#448)

* tle: move distributed APIs under language for v3.3.x

* tle: tolerate ImportError when optional raw module is unavailable

Author: sunnycase

python/triton/experimental/tle/__init__.py

@@ -1,43 +1,14 @@
 # flagtree tle
-from .distributed import (
-    B,
-    P,
-    S,
-    ShardedTensor,
-    ShardingSpec,
-    device_mesh,
-    distributed_barrier,
-    distributed_dot,
-    make_sharded_tensor,
-    remote,
-    reshard,
-    shard_id,
-    sharding,
-)
-
 from . import language
 
-# try:
-#     from . import raw
-# except ModuleNotFoundError:
-#     raw = None
+try:
+    from . import raw
+except (ModuleNotFoundError, ImportError):
+    raw = None
 
 __all__ = [
-    "device_mesh",
-    "S",
-    "P",
-    "B",
-    "sharding",
-    "ShardingSpec",
-    "ShardedTensor",
-    "make_sharded_tensor",
-    "reshard",
-    "remote",
-    "shard_id",
-    "distributed_barrier",
-    "distributed_dot",
     "language",
 ]
 
-# if raw is not None:
-#     __all__.append("raw")
+if raw is not None:
+    __all__.append("raw")

python/triton/experimental/tle/language/__init__.py

@@ -1,10 +1,39 @@
 # flagtree tle
 from .core import (
     load, )
+from .distributed import (
+    B,
+    P,
+    S,
+    ShardedTensor,
+    ShardingSpec,
+    device_mesh,
+    distributed_barrier,
+    distributed_dot,
+    make_sharded_tensor,
+    remote,
+    reshard,
+    shard_id,
+    sharding,
+)
 
 __all__ = [
     "load",
+    "device_mesh",
+    "S",
+    "P",
+    "B",
+    "sharding",
+    "ShardingSpec",
+    "ShardedTensor",
+    "make_sharded_tensor",
+    "reshard",
+    "remote",
+    "shard_id",
+    "distributed_barrier",
+    "distributed_dot",
+    "distributed",
     "dsa",
 ]
 
-from . import dsa
+from . import distributed, dsa

…n/triton/experimental/tle/distributed.py …experimental/tle/language/distributed.pypython/triton/experimental/tle/distributed.py renamed to python/triton/experimental/tle/language/distributed.py python/triton/experimental/tle/distributed.py renamed to python/triton/experimental/tle/language/distributed.py

@@ -711,7 +711,7 @@ def remote(
 
     Supported input:
     - tle buffered_tensor: returns a remote-marked buffered tensor; caller
-      should then use `tleg.local_ptr(...)` to materialize remote pointers.
+      should then use `tle.dsa.local_ptr(...)` to materialize remote pointers.
 
     `shard_id` is the target block id inside the current thread block cluster.
     When `scope` is provided, launch cluster dimensions are inferred from that
@@ -730,7 +730,7 @@ def remote(
         if (hasattr(tensor, "_tle_remote_shard_id") or hasattr(tensor, "_tle_remote_scope")
                 or hasattr(tensor.type, "_tle_remote_shard_id") or hasattr(tensor.type, "_tle_remote_scope")):
             raise ValueError("remote(buffered_tensor, ...) cannot be applied twice; "
-                             "materialize pointer views with tleg.local_ptr(remote_buffer, indices)")
+                             "materialize pointer views with tle.dsa.local_ptr(remote_buffer, indices)")
         if isinstance(shard_id, (int, tuple, list)):
             shard_id = _normalize_compile_time_remote_shard_id(shard_id, scope)
         else:

third_party/tsingmicro/examples/tle/test_tle_dsa_noc_gemm_4096.py

@@ -2,7 +2,7 @@
 import torch
 import triton
 import triton.language as tl
-from triton.experimental import tle
+import triton.experimental.tle.language as tle
 
 TILE_NUM = 16
 M = 4096
@@ -52,17 +52,17 @@ def dsa_shift_n_gemm_kernel(
     b_ptrs = B_ptr + offs_k[:, None] * N + offs_sub_n[None, :]
     b_init = tl.load(b_ptrs)
 
-    send_buf = tle.language.dsa.alloc((BLOCK_K, SUB_N), tl.float16)
-    recv_buf = tle.language.dsa.alloc((BLOCK_K, SUB_N), tl.float16)
+    send_buf = tle.dsa.alloc((BLOCK_K, SUB_N), tl.float16)
+    recv_buf = tle.dsa.alloc((BLOCK_K, SUB_N), tl.float16)
 
     offs_buf_k = tl.arange(0, BLOCK_K)[:, None] + tl.zeros((1, SUB_N), dtype=tl.int32)
     offs_buf_n = tl.arange(0, SUB_N)[None, :] + tl.zeros((BLOCK_K, 1), dtype=tl.int32)
 
-    send_ptr = tle.language.dsa.local_ptr(send_buf, [offs_buf_k, offs_buf_n])
-    recv_ptr = tle.language.dsa.local_ptr(recv_buf, [offs_buf_k, offs_buf_n])
+    send_ptr = tle.dsa.local_ptr(send_buf, [offs_buf_k, offs_buf_n])
+    recv_ptr = tle.dsa.local_ptr(recv_buf, [offs_buf_k, offs_buf_n])
 
     remote_recv_buf = tle.remote(recv_buf, send_next_tile)
-    remote_recv_ptr = tle.language.dsa.local_ptr(remote_recv_buf, [offs_buf_k, offs_buf_n])
+    remote_recv_ptr = tle.dsa.local_ptr(remote_recv_buf, [offs_buf_k, offs_buf_n])
 
     tl.store(send_ptr, b_init)