[Gluon] Add zeros, zeros_like and full_like (#7151)

Author: peterbell10

python/test/gluon/test_frontend.py

@@ -881,3 +881,30 @@ def test_tensor_reshape():
     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)
+
+
+@filecheck_test
+@gluon.jit
+def test_zeros():
+    # CHECK: [[BLOCKED:#.*]] = #ttg.blocked<{sizePerThread = [2]
+    # CHECK: [[BLOCKED2D:#.*]] = #ttg.blocked<{sizePerThread = [1, 2]
+    layout: ttgl.constexpr = ttgl.BlockedLayout([2], [32], [4], [0])
+    layout_2d: ttgl.constexpr = ttgl.BlockedLayout([1, 2], [4, 8], [4, 1], [1, 0])
+
+    # CHECK: arith.constant dense<0.000000e+00> : tensor<32xf32, [[BLOCKED]]>
+    a = ttgl.zeros([32], ttgl.float32, layout)
+
+    # CHECK: arith.constant dense<7.000000e+00> : tensor<32xf32, [[BLOCKED]]>
+    ttgl.full_like(a, 7)
+
+    # CHECK: arith.constant dense<0.000000e+00> : tensor<32xf32, [[BLOCKED]]>
+    ttgl.zeros_like(a)
+
+    # CHECK: arith.constant dense<0.000000e+00> : tensor<64xf32, [[BLOCKED]]>
+    ttgl.zeros_like(a, shape=[64])
+
+    # CHECK: arith.constant dense<0> : tensor<16x16xi8, [[BLOCKED2D]]>
+    ttgl.zeros_like(a, shape=[16, 16], dtype=ttgl.int8, layout=layout_2d)
+
+    # CHECK: arith.constant dense<7> : tensor<8x8xi16, [[BLOCKED2D]]>
+    ttgl.full_like(a, 7, shape=[8, 8], dtype=ttgl.int16, layout=layout_2d)

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

@@ -3,6 +3,7 @@
 import triton.language.standard as tl_standard
 from .._runtime import jit
 from triton import knobs
+from . import _core as ttgl
 
 _IMPORT_FROM_TRITON = [
     "sum",
@@ -12,10 +13,35 @@
     "xor_sum",
 ]
 
-__all__ = _IMPORT_FROM_TRITON
+__all__ = [
+    "full_like",
+    "zeros",
+    "zeros_like",
+    *_IMPORT_FROM_TRITON,
+]
 
 for name in _IMPORT_FROM_TRITON:
     # Convert JITFunction -> GluonJitFunction
     fn = getattr(tl_standard, name)
     assert knobs.runtime.interpret or isinstance(fn, triton.runtime.JITFunction)
     globals()[name] = jit(fn.fn)
+
+
+@jit
+def zeros(shape, dtype, layout):
+    return ttgl.full(shape, 0, dtype, layout)
+
+
+@jit
+def full_like(input, value, shape=None, dtype=None, layout=None):
+    return ttgl.full(
+        input.shape if shape is None else shape,
+        value,
+        input.dtype if dtype is None else dtype,
+        input.type.layout if layout is None else layout,
+    )
+
+
+@jit
+def zeros_like(input, shape=None, dtype=None, layout=None):
+    return full_like(input, 0, shape=shape, dtype=dtype, layout=layout)