partition working

python/tvm/contrib/cutlass/build.py

@@ -87,6 +87,9 @@ def visit_call(self, call):
         if str(op) == "nn.conv2d":
             self.op_attrs = call.attrs
 
+        for arg in call.args:
+            self.visit(arg)
+
 
 def select_gemm_kernel(
     cutlass_profiler, MM, KK, NN, out_dtype, batched, profile_all, use_multiprocessing

python/tvm/relay/op/contrib/cutlass.py

@@ -56,7 +56,7 @@ def make_batch_matmul_pattern():
     return is_op("nn.batch_matmul")(wildcard(), wildcard())
 
 
-def make_conv2d_pattern(with_bias=True, with_act=None, out_dtype="float16"):
+def make_conv2d_pattern(with_bias=False, with_act=None, out_dtype="float16"):
     """Create a pattern for dense op followed by activations."""
     data = wildcard()
     weight = wildcard()

tests/python/contrib/test_cutlass.py

@@ -385,16 +385,24 @@ def test_conv2d():
     )
 
 
+d_shape = (16, 16, 32, 32)
+w_shape = (32, 16, 3, 3)
+padding = (1, 1)
+
+
 def test_conv2d_bias():
-    d_shape = (16, 16, 32, 32)
-    w_shape = (32, 16, 3, 3)
-    padding = (1, 1)
     mod_nchw = get_conv2d_nchw_bias(d_shape, w_shape, padding)
+    verify_conv2d(
+        mod_nchw, mod_nchw, d_shape, w_shape, sm=80, atol=1e-5, rtol=1e-5, run_benchmark=False
+    )
+
 
+def test_conv2d_bias_relu():
+    mod_nchw = get_conv2d_nchw_bias_relu(d_shape, w_shape, padding)
     verify_conv2d(
         mod_nchw, mod_nchw, d_shape, w_shape, sm=80, atol=1e-5, rtol=1e-5, run_benchmark=False
     )
 
 
 if __name__ == "__main__":
-    test_conv2d_bias()
+    test_conv2d_bias_relu()