test worked with fp32 output

masahi · Feb 4, 2022 · 08a6147 · 08a6147
1 parent 084d5c4
commit 08a6147
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 9 deletions.
diff --git a/python/tvm/relay/op/nn/_nn.py b/python/tvm/relay/op/nn/_nn.py
@@ -1107,6 +1107,7 @@ def legalize_conv2d_backward_weight(attrs, inputs, types):
         padding=attrs.padding,
         dilation=attrs.strides,
         groups=in_channel * batch,
+        out_dtype=attrs.out_dtype
     )
 
     # infer shape of backward_weight

diff --git a/src/relay/backend/contrib/cutlass/codegen.cc b/src/relay/backend/contrib/cutlass/codegen.cc
@@ -319,7 +319,7 @@ std::string Conv2dOp(std::string id, const Str2StrMap& attrs,
   CutlassPrint(conv2d_decl, "int dilation_h = " + attrs.at("dilation_h") + ";\n");
   CutlassPrint(conv2d_decl, "int dilation_w = " + attrs.at("dilation_w") + ";\n");
   // TODO
-  const int split_k_slices = 1;
+  const int split_k_slices = 8;
   CutlassPrint(conv2d_decl, "int split_k_slices = " + std::to_string(split_k_slices) + ";\n");
 
   CutlassPrint(
@@ -329,7 +329,7 @@ std::string Conv2dOp(std::string id, const Str2StrMap& attrs,
       "split_k_slices);\n");
 
   const bool use_split_k = split_k_slices > 1;
-  const std::string split_k_mode = use_split_k > 1 ? "kParallel" : "kSerial";
+  const std::string split_k_mode = use_split_k ? "kParallel" : "kSerial";
   CutlassPrint(conv2d_decl,
                "const cutlass::conv::SplitKMode split_k_mode = cutlass::conv::SplitKMode::" +
                    split_k_mode + ";\n");

diff --git a/src/relay/op/nn/convolution.cc b/src/relay/op/nn/convolution.cc
@@ -643,7 +643,7 @@ bool Conv2DBackwardWeightRel(const Array<Type>& types, int num_inputs, const Att
       {out_channels, in_channels, param->kernel_size[0], param->kernel_size[1]});
 
   auto wshape = trans_kernel_layout.BackwardShape(wshape_oihw);
-  reporter->Assign(types[2], TensorType(wshape, data->dtype));
+  reporter->Assign(types[2], TensorType(wshape, param->out_dtype));
   return true;
 }
 

diff --git a/tests/python/contrib/test_cutlass.py b/tests/python/contrib/test_cutlass.py
@@ -829,11 +829,11 @@ def test_conv2d_transpose():
 
 
 def test_conv2d_backward_weight():
-    OC = 8
-    IC = 16
+    OC = 32
+    IC = 32
     d_shape = (16, IC, 32, 32)
     w_shape = (OC, IC, 3, 3)
-    dtype = "float32"
+    dtype = "float16"
 
     for strides in [(1, 1), (2, 2)]:
         o_shape = (16, OC, 32 // strides[0], 32 // strides[1])
@@ -845,7 +845,7 @@ def test_conv2d_backward_weight():
             o_shape,
             padding,
             strides,
-            out_dtype=dtype,
+            out_dtype="float32",
             data_dtype=dtype,
             weight_dtype=dtype,
         )
@@ -863,8 +863,6 @@ def test_conv2d_backward_weight():
             data_dtype=dtype,
         )
 
-    # split k
-
 
 def test_conv2d_bwd():
     IC = 16