[FlagGems Operator Development Competition] Optimize roll

Author: HyosungSink

benchmark/core_shapes.yaml

@@ -63,6 +63,15 @@ log10_out:
     - [1024, 65536]
   shape_desc: "(B), M, N"
 
+roll:
+  shapes:
+    - [1024]
+    - [256, 256]
+    - [32, 64, 64]
+    - [64, 128, 128]
+    - [8, 16, 128, 128]
+  shape_desc: "layout roll representative shapes"
+
 softmax_backward:
   shapes:
     - [1048576] # 1024 * 1024

benchmark/test_generic_pointwise_perf.py

@@ -23,7 +23,10 @@ def flip_input_fn(shape, cur_dtype, device):
 
 def roll_input_fn(shape, cur_dtype, device):
     inp = generate_tensor_input(shape, cur_dtype, device)
+    yield inp, {"shifts": 257}
     if len(shape) > 1:
+        yield inp, {"shifts": 7, "dims": 0}
+        yield inp, {"shifts": 13, "dims": len(shape) - 1}
         yield inp, {"shifts": (1, 2), "dims": (0, 1)}
     else:
         yield inp, {"shifts": 1, "dims": 0}