Conv Optimization && Fix

operatorspy/tests/conv.py

@@ -41,14 +41,32 @@ class ConvDescriptor(Structure):
 def conv(x, w, stride, padding, dilation):
     match len(x.shape) - 2:
         case 1:
+            if PROFILE:
+                ans = F.conv1d(
+                    x, w, stride=stride, padding=padding, dilation=dilation
+                )
+                torch.cuda.synchronize()
+                return ans
             return F.conv1d(
                 x, w, stride=stride, padding=padding, dilation=dilation
             )
         case 2:
+            if PROFILE:
+                ans = F.conv2d(
+                        x, w, stride=stride, padding=padding, dilation=dilation
+                    )
+                torch.cuda.synchronize()
+                return ans
             return F.conv2d(
                 x, w, stride=stride, padding=padding, dilation=dilation
             )
         case 3:
+            if PROFILE:
+                ans = F.conv3d(
+                        x, w, stride=stride, padding=padding, dilation=dilation
+                    )
+                torch.cuda.synchronize()
+                return ans
             return F.conv3d(
                 x, w, stride=stride, padding=padding, dilation=dilation
             )

src/ops/conv/cuda/conv.cc

@@ -92,6 +92,7 @@ infiniopStatus_t cudaCreateConvDescriptor(CudaHandle_t handle,
     checkCudnnError(cudnnCreateTensorDescriptor(&y_desc));
     checkCudnnError(cudnnSetTensorNdDescriptorEx(y_desc, CUDNN_TENSOR_NCHW, static_cast<cudnnDataType_t>(tensor_dt), new_ndim, y_shape));
 
+    cudnnSetConvolutionMathType(op_desc, CUDNN_TENSOR_OP_MATH_ALLOW_CONVERSION);
 
     // tuning: get the best algorithm
     int requestedAlgoCount = 1;