[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

@@ -158,6 +158,7 @@ def enable(lib=aten_lib):
     lib.impl("repeat_interleave.self_int", repeat_interleave_self_int, "CUDA")
     lib.impl("vstack", vstack, "CUDA")
     lib.impl("repeat_interleave.Tensor", repeat_interleave_tensor, "CUDA")
+    lib.impl("randperm", randperm, "CUDA")
 
 
 class use_gems:

src/flag_gems/ops/__init__.py

@@ -76,6 +76,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 .repeat import repeat
@@ -185,6 +186,7 @@
     "minimum",
     "rand",
     "randn",
+    "randperm",
     "rand_like",
     "randn_like",
     "resolve_neg",