2222from  tvm  import  relay 
2323from  tvm .relay  import  transform 
2424from  tvm .relay .testing  import  ctx_list 
25+ from  tvm .contrib  import  util 
2526import  topi .testing 
2627
2728def  run_infer_type (expr ):
@@ -134,6 +135,46 @@ def run_test_conv2d(dtype, out_dtype, scale, dshape, kshape,
134135            op_res1  =  intrp1 .evaluate (func )(data , kernel )
135136            tvm .testing .assert_allclose (op_res1 .asnumpy (), ref_res , rtol = 1e-5 , atol = 1e-5 )
136137
138+     def  compile_test_conv2d_arm_cpu (dtype , out_dtype , scale , dshape , kshape ,
139+                         padding = (1 , 1 ),
140+                         groups = 1 ,
141+                         dilation = (1 , 1 ),
142+                         ** attrs ):
143+         x  =  relay .var ("x" , shape = dshape , dtype = dtype )
144+         w  =  relay .var ("w" , dtype = dtype )
145+         y  =  relay .nn .conv2d (x , w ,
146+                             padding = padding ,
147+                             dilation = dilation ,
148+                             groups = groups ,
149+                             ** attrs )
150+         func  =  relay .Function ([x , w ], y )
151+         mod  =  tvm .relay .Module ()
152+         mod ["main" ] =  func 
153+ 
154+         test_schedule = '{"i": ["llvm -device=arm_cpu", "topi_nn_depthwise_conv2d_nchw", \  
155+ \ 
156+ \ 
157+ \ 
158+ \ 
159+ \ 
160+ \ 
161+ \ 
162+ \ 
163+ \ 
164+ \ 
165+ \ 
166+ \ 
167+ \ 
168+ \ 
169+ 
170+         temp  =  util .tempdir ()
171+         with  open (temp .relpath ("temp.log" ), "w" ) as  log_file :
172+             log_file .write (test_schedule )
173+         with  autotvm .apply_history_best (temp .relpath ("temp.log" )):
174+             with  relay .build_config (opt_level = 3 ):
175+                 print ('Compiling...' )
176+                 graph_json , mod , params  =  tvm .relay .build (mod , target = "llvm -device=arm_cpu" )
177+ 
137178    # depthwise conv2d 
138179    dshape  =  (1 , 32 , 18 , 18 )
139180    kshape  =  (32 , 1 , 3 , 3 )
@@ -142,6 +183,13 @@ def run_test_conv2d(dtype, out_dtype, scale, dshape, kshape,
142183                    fref = lambda  x , w : topi .testing .depthwise_conv2d_python_nchw (
143184                        x , w , (1 , 1 ), "SAME" ))
144185
186+     # depthwise conv2d for arm_cpu 
187+     dshape  =  (1 , 512 , 32 , 32 )
188+     kshape  =  (512 , 1 , 3 , 3 )
189+     compile_test_conv2d_arm_cpu ("float32" , "float32" , 1 , dshape , kshape ,
190+                                 padding = (1 , 1 ), channels = 512 , 
191+                                 groups = 512 , kernel_size = (3  ,3 ))
192+ 
145193    # CUDA is disabled for 'direct' schedule: 
146194    # https://github.com/apache/incubator-tvm/pull/3070#issuecomment-486597553 
147195    # group conv2d 
0 commit comments