[Operator] Add tile op (#148)

Co-authored-by: Clement Chan <[email protected]>

benchmark/test_pointwise_perf.py

@@ -614,3 +614,19 @@ def masked_fill_args(dtype, batch, size):
         sizes=SIZES,
     )
     bench.run()
+
+
+def test_perf_tile():
+    def tile_kwargs(dtype, batch, size):
+        return {"dims": [2, 4]}
+
+    bench = Benchmark(
+        op_name="tile",
+        torch_op=torch.tile,
+        arg_func=unary_arg,
+        dtypes=FLOAT_DTYPES,
+        batch=POINTWISE_BATCH,
+        sizes=SIZES,
+        kwargs_func=tile_kwargs,
+    )
+    bench.run()

src/flag_gems/__init__.py

@@ -126,6 +126,7 @@ def enable(lib=aten_lib):
     lib.impl("isclose", isclose, "CUDA")
     lib.impl("allclose", allclose, "CUDA")
     lib.impl("flip", flip, "CUDA")
+    lib.impl("tile", tile, "CUDA")
     lib.impl("index_select", index_select, "CUDA")
     lib.impl("masked_fill", masked_fill, "CUDA")
     lib.impl("_unique2", _unique2, "CUDA")

src/flag_gems/ops/__init__.py

@@ -79,6 +79,7 @@
 from .sub import sub
 from .sum import sum, sum_dim
 from .tanh import tanh
+from .tile import tile
 from .topk import topk
 from .triu import triu
 from .uniform import uniform_
@@ -179,6 +180,7 @@
     "softmax",
     "sub",
     "tanh",
+    "tile",
     "triu",
     "topk",
     "max",