Add conv ops to operator microbenchmark (pytorch#166331)

Adding `conv` (conv1d, conv2d, conv3d) to the list of operator microbenchmarks run in the CI script (`.ci/pytorch/test.sh`), ensuring convolution operators are now benchmarked alongside existing ones.
Pull Request resolved: pytorch#166331
Approved by: https://github.com/huydhn, https://github.com/jbschlosser

Author: jainapurva

.ci/pytorch/test.sh

@@ -1653,7 +1653,7 @@ test_operator_microbenchmark() {
 
   cd "${TEST_DIR}"/benchmarks/operator_benchmark
 
-  for OP_BENCHMARK_TESTS in matmul mm addmm bmm; do
+  for OP_BENCHMARK_TESTS in matmul mm addmm bmm conv; do
     $TASKSET python -m pt.${OP_BENCHMARK_TESTS}_test --tag-filter long \
       --output-json-for-dashboard "${TEST_REPORTS_DIR}/operator_microbenchmark_${OP_BENCHMARK_TESTS}_compile.json" \
       --benchmark-name "PyTorch operator microbenchmark" --use-compile

benchmarks/operator_benchmark/pt/configs.py

@@ -11,6 +11,11 @@ def remove_cuda(config_list):
     return [config for config in config_list if cuda_config not in config]
 
 
+def remove_cpu(config_list):
+    cpu_config = {"device": "cpu"}
+    return [config for config in config_list if cpu_config not in config]
+
+
 # Configs for conv-1d ops
 conv_1d_configs_short = op_bench.config_list(
     attr_names=["IC", "OC", "kernel", "stride", "N", "L"],
@@ -127,6 +132,18 @@ def remove_cuda(config_list):
     },
     tags=["short"],
 )
+conv_3d_configs_long = op_bench.cross_product_configs(
+    IC=[16, 32],
+    OC=[32, 64],
+    kernel=[3, 5],
+    stride=[1, 2],
+    N=[1],
+    D=[128],
+    H=[128],
+    W=[128],
+    device=["cpu", "cuda"],
+    tags=["long"],
+)
 
 linear_configs_short = op_bench.config_list(
     attr_names=["N", "IN", "OUT"],

benchmarks/operator_benchmark/pt/conv_test.py

@@ -38,6 +38,10 @@ def forward(self, input):
 op_bench.generate_pt_test(
     configs.conv_1d_configs_short + configs.conv_1d_configs_long, Conv1dBenchmark
 )
+op_bench.generate_pt_gradient_test(
+    configs.remove_cpu(configs.conv_1d_configs_short + configs.conv_1d_configs_long),
+    Conv1dBenchmark,
+)
 
 
 if not torch.backends.mkldnn.is_acl_available():
@@ -103,6 +107,20 @@ def forward(self, input):
     configs.conv_2d_pw_configs_short + configs.conv_2d_pw_configs_long,
     Conv2dPointwiseBenchmark,
 )
+op_bench.generate_pt_gradient_test(
+    configs.remove_cpu(configs.conv_2d_configs_short + configs.conv_2d_configs_long),
+    Conv2dBenchmark,
+)
+op_bench.generate_pt_gradient_test(
+    configs.remove_cpu(configs.conv_2d_configs_short + configs.conv_2d_configs_long),
+    ConvTranspose2dBenchmark,
+)
+op_bench.generate_pt_gradient_test(
+    configs.remove_cpu(
+        configs.conv_2d_pw_configs_short + configs.conv_2d_pw_configs_long
+    ),
+    Conv2dPointwiseBenchmark,
+)
 
 
 """
@@ -134,6 +152,12 @@ def forward(self, input):
 
 op_bench.generate_pt_test(configs.conv_3d_configs_short, Conv3dBenchmark)
 op_bench.generate_pt_test(configs.conv_3d_configs_short, ConvTranspose3dBenchmark)
+op_bench.generate_pt_gradient_test(
+    configs.remove_cpu(configs.conv_3d_configs_long), Conv3dBenchmark
+)
+op_bench.generate_pt_gradient_test(
+    configs.remove_cpu(configs.conv_3d_configs_long), ConvTranspose3dBenchmark
+)
 
 
 if __name__ == "__main__":