[Operator] Add randperm op

benchmark/test_tensor_constructor_perf.py

@@ -159,3 +159,19 @@ def full_kwargs(dtype, batch, size):
         kwargs_func=full_kwargs,
     )
     bench.run()
+
+
+def test_perf_randperm():
+    def randperm_args(dtype, batch, size):
+        return {"n": size, "dtype": dtype, "device": "cuda"}
+
+    bench = Benchmark(
+        op_name="randperm",
+        torch_op=torch.randperm,
+        arg_func=None,
+        dtypes=[torch.int32, torch.int64],
+        batch=POINTWISE_BATCH,
+        sizes=SIZES,
+        kwargs_func=randperm_args,
+    )
+    bench.run()

src/flag_gems/__init__.py

@@ -131,6 +131,7 @@ def enable(lib=aten_lib):
     lib.impl("index_select", index_select, "CUDA")
     lib.impl("masked_fill", masked_fill, "CUDA")
     lib.impl("_unique2", _unique2, "CUDA")
+    lib.impl("randperm", randperm, "CUDA")
 
 
 class use_gems:

src/flag_gems/ops/__init__.py

@@ -67,6 +67,7 @@
 from .rand_like import rand_like
 from .randn import randn
 from .randn_like import randn_like
+from .randperm import randperm
 from .reciprocal import reciprocal
 from .relu import relu
 from .resolve_conj import resolve_conj
@@ -157,6 +158,7 @@
     "minimum",
     "rand",
     "randn",
+    "randperm",
     "rand_like",
     "randn_like",
     "resolve_neg",