Add support Conv and pooling

Author: spike-zhu

operatorspy/tests/avg_pool.py

@@ -152,10 +152,6 @@ def test(
         elapsed = (time.time() - start_time) / NUM_ITERATIONS
         print(f"    lib time: {elapsed :6f}")
 
-
-    print(x)
-    print(y)
-    print(ans)
     assert torch.allclose(y, ans, atol=0, rtol=1e-3)
     check_error(lib.infiniopDestroyAvgPoolDescriptor(descriptor))
 
@@ -205,7 +201,7 @@ def test_musa(lib, test_cases):
         # ((1, 1, 10), (3,), (1,), (1,)),
         ((1, 1, 2, 2), (2, 2), (1, 1), (1, 1)),
         ((32, 4, 224, 224), (3, 3), (1, 1), (2, 2)),
-        ((1, 1, 16, 16, 16), (5, 5, 5), (2, 2, 2), (2, 2, 2)),
+        # ((1, 1, 16, 16, 16), (5, 5, 5), (2, 2, 2), (2, 2, 2)),
     ]
     args = get_args()
     lib = open_lib()

operatorspy/tests/global_avg_pool.py

@@ -158,7 +158,8 @@ def test_musa(lib, test_cases):
     test_cases = [
         # x_shape
         ((1, 1, 2, 2)),
-        # ((1, 3, 3)),
+        ((1, 3, 3, 6)),
+        ((8, 4, 6, 12)),
         # ((1, 3, 1, 1, 3)),
         # ((1, 3, 1, 1, 257)),
         # ((1, 2, 1, 1, 514)),

operatorspy/tests/max_pool.py

@@ -148,9 +148,9 @@ def test(
             )
         elapsed = (time.time() - start_time) / NUM_ITERATIONS
         print(f"    lib time: {elapsed :6f}")
-    print(x)
-    print(y)
-    print(ans)
+    # print(x)
+    # print(y)
+    # print(ans)
     assert torch.allclose(y, ans, atol=0, rtol=1e-3)
     check_error(lib.infiniopDestroyMaxPoolDescriptor(descriptor))
 
@@ -198,7 +198,9 @@ def test_musa(lib, test_cases):
     test_cases = [
         # x_shape, kernel_shape, padding, strides
         # ((1, 1, 10), (3,), (1,), (1,)),
-        # ((32, 3, 224, 224), (3, 3), (1, 1), (2, 2)),
+        ((32, 3, 224, 224), (3, 3), (1, 1), (2, 2)),
+        ((1, 3, 6, 6), (3, 3), (1, 1), (2, 2)),
+        ((8, 3, 12, 12), (3, 3), (1, 1), (2, 2)),
         ((1, 1, 4, 4), (2, 2), (0, 0), (1, 1)),
         # ((1, 1, 16, 16, 16), (5, 5, 5), (2, 2, 2), (2, 2, 2)),
     ]

src/ops/add/musa/add_musa.mu

@@ -84,6 +84,7 @@ void _add_mt_gpu(AddMusaDescriptor_t desc, Tdata *c, Tdata const *a, Tdata const
         add<Tdata, BTdata><<<gridDims, blockDims, 0, musa_stream>>>(
             c, a, b, desc->a_strides, desc->b_strides, desc->c_strides, offset + data_size, desc->ndim, offset + i, desc->broadcasted, pack_size);
     }
+    printf("[SUCCESS to execute add_mt_gpu]\n");
 }
 
 template<typename Tdata, typename TIdata>

src/ops/pooling/musa/pooling_musa.h

@@ -46,6 +46,7 @@ inline musa::dnn::Pooling::Mode getPoolingMode(int pooling_type) {
             return musa::dnn::Pooling::Mode::MAXPOOL;
         case 1:
             return musa::dnn::Pooling::Mode::AVGPOOL_COUNT_PAD;
+            // return musa::dnn::Pooling::Mode::AVGPOOL_COUNT_WITHOUT_PAD;
         default:
             return musa::dnn::Pooling::Mode::MAXPOOL;
     }