[FRONTEND] Refactor unsplat to use new op (#7586)

Co-authored-by: peterbell10 <[email protected]>

Author: ThomasRaoux

lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp

@@ -529,6 +529,7 @@ void populateTritonPatterns(TritonGPUTypeConverter &typeConverter,
       GenericOpPattern<triton::IntToPtrOp>,
       GenericOpPattern<triton::PtrToIntOp>,
       GenericOpPattern<triton::SplatOp>,
+      GenericOpPattern<triton::UnsplatOp>,
       GenericOpPattern<triton::AddPtrOp>,
       TritonBroadcastPattern,
       TritonCatPattern,

python/src/ir.cc

@@ -1495,6 +1495,10 @@ void init_triton_ir(py::module &&m) {
            [](TritonOpBuilder &self, Type &retTy, Value &arg) -> Value {
              return self.createOrFold<SplatOp>(retTy, arg);
            })
+      .def("create_unsplat",
+           [](TritonOpBuilder &self, Value &arg) -> Value {
+             return self.createOrFold<UnsplatOp>(arg);
+           })
       // // atomic
       .def("create_atomic_cas",
            [](TritonOpBuilder &self, Value &ptr, Value &cmp, Value &val,

python/triton/language/core.py

@@ -10,7 +10,7 @@
 from dataclasses import dataclass
 import builtins
 from .. import knobs
-from ..runtime.jit import jit, JITFunction
+from ..runtime.jit import JITFunction
 import inspect
 
 from .._C.libtriton import ir
@@ -1831,11 +1831,6 @@ def join(a, b, _semantic=None):
     return _semantic.join(a, b)
 
 
-@jit
-def _take_first(a, b):
-    return a
-
-
 def _unsplat(x, _semantic=None, _generator=None):
     """
     Convert a single-element tensor to a scalar.
@@ -1846,10 +1841,7 @@ def _unsplat(x, _semantic=None, _generator=None):
     for d in x.shape:
         numel *= d
     assert numel == 1, "can only unsplat single-element tensors"
-    if len(x.shape) >= 2:
-        x = _semantic.reshape(x, [1])
-    x = typing.cast(tensor, reduce(x, 0, _take_first, _semantic=_semantic, _generator=_generator))
-    return x
+    return _semantic.unsplat(x)
 
 
 @_tensor_member_fn

python/triton/language/semantic.py

@@ -619,6 +619,9 @@ def splat(self, value: TensorTy, shape: List[int]) -> TensorTy:
         ret_ty = tl.block_type(value.dtype, shape)
         return self.tensor(self.builder.create_splat(ret_ty.to_ir(self.builder), value.handle), ret_ty)
 
+    def unsplat(self, value: TensorTy) -> TensorTy:
+        return self.tensor(self.builder.create_unsplat(value.handle), value.dtype)
+
     def reshape(self, input: TensorTy, dst_shape: List[int], can_reorder: bool) -> TensorTy:
         numel = 1
         for s in dst_shape:

python/triton/runtime/interpreter.py

@@ -663,6 +663,9 @@ def create_splat(self, ret_ty, arg):
         else:  # scalar
             return TensorHandle(np.full(shape, arg.data, dtype=_get_np_dtype(arg.dtype)), arg.dtype.scalar)
 
+    def create_unsplat(self, arg):
+        return TensorHandle(np.full((1, ), arg.data[0], dtype=_get_np_dtype(arg.dtype)), arg.dtype.scalar)
+
     def create_atomic_cas(self, ptr, cmp, val, sem, scope):
         if sem not in self.ir_sem_to_interpreter_sem:
             raise ValueError(f"unsupported semantic {sem}")