Add cudnn_convolution benchmark

Author: Yangyang0906C

benchmark/test_convolution_perf.py

@@ -173,6 +173,62 @@ def conv_depthwise2d_input_fn(shape, dtype, device):
     bench.run()
 
 
+class CudnnConvolutionBenchmark(GenericBenchmark):
+    def set_more_shapes(self):
+        return [
+            (1, 2, 5, 5, 1, 3, 3, 1, 0, 1, 1),
+            (2, 3, 9, 9, 1, 3, 3, 1, 1, 1, 1),
+            (32, 8, 8, 8, 32, 2, 2, 1, 0, 1, 1),
+            (32, 64, 128, 128, 32, 3, 3, 1, 2, 1, 1),
+            (32, 64, 210, 210, 16, 5, 5, 2, 1, 1, 1),
+            (16, 32, 24, 24, 24, 3, 3, 2, 2, 2, 2),
+        ]
+
+
+@pytest.mark.cudnn_convolution
+def test_perf_cudnn_convolution():
+    def cudnn_conv_input_fn(shape, dtype, device):
+        (
+            batch,
+            input_c,
+            input_h,
+            input_w,
+            out_c,
+            kernel_h,
+            kernel_w,
+            stride,
+            padding,
+            groups,
+            dilation,
+        ) = shape
+        input_shape = (batch, input_c, input_h, input_w)
+        weight_shape = (out_c, input_c // groups, kernel_h, kernel_w)
+        input = torch.randn(size=input_shape, device=device, dtype=dtype)
+        weight = torch.randn(size=weight_shape, device=device, dtype=dtype)
+
+        yield {
+            "input": input,
+            "weight": weight,
+            "padding": [padding, padding],
+            "stride": [stride, stride],
+            "dilation": [dilation, dilation],
+            "groups": groups,
+            "benchmark": False,
+            "deterministic": False,
+            "allow_tf32": False,
+        },
+
+    torch.backends.cudnn.allow_tf32 = False
+    bench = CudnnConvolutionBenchmark(
+        input_fn=cudnn_conv_input_fn,
+        op_name="cudnn_convolution",
+        torch_op=torch.cudnn_convolution,
+        dtypes=[torch.float16, torch.float32],
+    )
+    bench.set_gems(flag_gems.cudnn_convolution)
+    bench.run()
+
+
 class Conv3DBenchmark(GenericBenchmark):
     def set_more_shapes(self):
         return None