[Gluon] Add split, join and reshape on tensor (#7122)

Author: peterbell10

python/test/gluon/test_frontend.py

@@ -844,3 +844,36 @@ def test_tensor_permute():
     res = ttgl.permute(a, [1, 0])
     permuted_layout: ttgl.constexpr = ttgl.BlockedLayout([2, 1], [8, 4], [1, 4], [0, 1], [1, 1], [1, 1], [1, 0])
     ttgl.static_assert(permuted_layout == res.type.layout)
+
+
+@filecheck_test
+@gluon.jit
+def test_split_join():
+    # CHECK: [[BLOCKED:#.*]] = #ttg.blocked<{sizePerThread = [2], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}>
+    # CHECK: [[BLOCKED1:#.*]] = #ttg.blocked<{sizePerThread = [2, 2], threadsPerWarp = [32, 1], warpsPerCTA = [4, 1], order = [1, 0]}>
+    layout: ttgl.constexpr = ttgl.BlockedLayout([2], [32], [4], [0])
+    a = ttgl.full([128], 1, ttgl.int32, layout)
+    b = ttgl.full([128], 2, ttgl.int32, layout)
+    # CHECK: tt.join {{.*}} : tensor<128xi32, [[BLOCKED]]> -> tensor<128x2xi32, [[BLOCKED1]]>
+    res = ttgl.join(a, b)
+    expect_layout: ttgl.constexpr = ttgl.BlockedLayout([2, 2], [32, 1], [4, 1], [1, 0], [1, 1], [1, 1], [1, 0])
+    ttgl.static_assert(res.type.layout == expect_layout)
+
+    # CHECK: tt.split {{.*}} : tensor<128x2xi32, [[BLOCKED1]]> -> tensor<128xi32, [[BLOCKED]]>
+    c, d = ttgl.split(res)
+    ttgl.static_assert(c.type.layout == layout)
+    ttgl.static_assert(d.type.layout == layout)
+
+
+@filecheck_test
+@gluon.jit
+def test_tensor_reshape():
+    # CHECK: [[BLOCKED:#.*]] = #ttg.blocked<{sizePerThread = [2], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}>
+    # CHECK: [[BLOCKED1:#.*]] = #ttg.blocked<{sizePerThread = [1, 1, 2], threadsPerWarp = [2, 4, 4], warpsPerCTA = [4, 1, 1], order = [2, 1, 0]}>
+    layout: ttgl.constexpr = ttgl.BlockedLayout([2], [32], [4], [0])
+    a = ttgl.full([256], 1, ttgl.int32, layout)
+    # CHECK: tt.reshape {{.*}} : tensor<256xi32, [[BLOCKED]]> -> tensor<8x4x8xi32, [[BLOCKED1]]>
+    v = a.reshape([8, 4, 8])
+    expect_layout: ttgl.constexpr = ttgl.BlockedLayout([1, 1, 2], [2, 4, 4], [4, 1, 1], [2, 1, 0], [1, 1, 1], [1, 1, 1],
+                                                       [2, 1, 0])
+    ttgl.static_assert(v.type.layout == expect_layout)

python/triton/experimental/gluon/language/_core.py

@@ -42,17 +42,20 @@
 )
 
 _IMPORT_FROM_TRITON: List[str] = [
-    "expand_dims",  # NOQA: F822
-    "load",  # NOQA: F822
-    "program_id",  # NOQA: F822
-    "reduce",  # NOQA: F822
-    "static_assert",  # NOQA: F822
-    "store",  # NOQA: F822
-    "to_tensor",  # NOQA: F822
-    "where",  # NOQA: F822
-    "maximum",  # NOQA: F822
-    "minimum",  # NOQA: F822
+    "expand_dims",
+    "join",
+    "load",
+    "maximum",
+    "minimum",
     "permute",
+    "program_id",
+    "reduce",
+    "reshape",
+    "split",
+    "static_assert",
+    "store",
+    "to_tensor",
+    "where",
 ]
 
 __all__ = [

python/triton/experimental/gluon/language/_semantic.py

@@ -22,6 +22,11 @@ class GluonSemantic(TritonSemantic[TensorTy]):
     def __init__(self, builder: GluonOpBuilder):
         self.builder = builder
 
+    def _wrap_tensor_infer_layout(self, tensor):
+        ty = ttgl.distributed_type(tensor.type.scalar, tensor.shape,
+                                   self.builder.get_gluon_layout_from_tensor(tensor.handle))
+        return self.tensor(tensor.handle, ty)
+
     def _broadcast_shapes(self, lhs_shape: List[int], rhs_shape: List[int]):
         if len(lhs_shape) != len(rhs_shape):
             raise ValueError(f"Cannot broadcast, rank mismatch: {lhs_shape}, {rhs_shape}")
@@ -57,11 +62,19 @@ def expand_dims(self, input: TensorTy, axis: int) -> TensorTy:
         handle = self.builder.create_expand_dims(input.handle, axis, ret_ty.to_ir(self.builder))
         return self.tensor(handle, ret_ty)
 
+    def join(self, a: TensorTy, b: TensorTy) -> TensorTy:
+        a, b = self.broadcast_impl_value(a, b)
+        _check(a.shape != [], "Cannot join scalars in gluon")
+        value = super().join(a, b)
+        return self._wrap_tensor_infer_layout(value)
+
+    def split(self, a: TensorTy) -> Tuple[TensorTy, TensorTy]:
+        lhs, rhs = super().split(a)
+        return self._wrap_tensor_infer_layout(lhs), self._wrap_tensor_infer_layout(rhs)
+
     def permute(self, input: TensorTy, dims: Tuple[int]) -> TensorTy:
         value = super().permute(input, dims)
-        layout = self.builder.get_gluon_layout_from_tensor(value.handle)
-        res_ty = ttgl.distributed_type(value.type.scalar, value.shape, layout)
-        return self.tensor(value.handle, res_ty)
+        return self._wrap_tensor_infer_layout(value)
 
     def broadcast_impl_shape(self, input: TensorTy, shape: Tuple[int]) -> TensorTy:
         _check(isinstance(input.type, ttgl.distributed_type),
@@ -106,6 +119,11 @@ def arange(self, start, end, layout):
         ret_ty = ttgl.distributed_type(ttgl.int32, shape, layout)
         return super().arange(start, end, ret_ty=ret_ty)
 
+    def reshape(self, input: TensorTy, dst_shape: List[int], can_reorder: bool):
+        _check(not can_reorder, "can_reorder is not supported in gluon")
+        value = super().reshape(input, dst_shape, can_reorder)
+        return self._wrap_tensor_infer_layout(value)
+
     def splat(self, value, shape, layout):
         ret_ty = ttgl.distributed_type(value.dtype, shape, layout)
         handle = self.builder.create_splat(ret_ty.to_ir(self.builder), value.handle)