Skip to content

Commit 517c420

Browse files
krishnaraj36Sanjay Shankar Krishnaa
andauthored
[TOPI][ADRENO] Add Group Conv2d texture schedule (#17274)
* Added Support for Adreno Texture Based Group Convolution * Added Few Testcases and Fixed Compute * Limited Support for Group Convolution * Removed Dead Code, Fixed Minor Issues --------- Co-authored-by: Sanjay Shankar Krishnaa <[email protected]>
1 parent 4a37f64 commit 517c420

File tree

4 files changed

+625
-1
lines changed

4 files changed

+625
-1
lines changed

python/tvm/relay/op/strategy/adreno.py

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -182,8 +182,37 @@ def conv2d_strategy_adreno(attrs, inputs, out_type, target):
182182
+ kernel_layout
183183
+ ") - only support NCHW4c / OIHW4o and NHWC / HWOI layouts for conv2d"
184184
)
185+
elif (data_layout == "NCHW4c" or data_layout == "NCHW") and (
186+
kernel_layout == "OIHW" or kernel_layout == "OIHW4o"
187+
):
188+
pad_in_chunks = (len(data.shape) == 5 and data.shape[1] % groups != 0) or (
189+
len(data.shape) == 4 and data.shape[1] % (groups * 4) != 0
190+
)
191+
pad_out_chunks = (len(kernel.shape) == 5 and kernel.shape[0] % groups != 0) or (
192+
len(kernel.shape) == 4 and kernel.shape[0] % (groups * 4) != 0
193+
)
194+
195+
if not (pad_in_chunks or pad_out_chunks):
196+
strategy.add_implementation(
197+
wrap_compute_conv2d(topi.adreno.group_conv2d_nchwc),
198+
wrap_topi_schedule(topi.adreno.schedule_group_conv2d_nchwc),
199+
name="group_conv2d_nchwc.image2d",
200+
plevel=10,
201+
)
202+
elif len(data.shape) == 4 and len(kernel.shape) == 4:
203+
strategy.add_implementation(
204+
wrap_compute_conv2d(topi.cuda.group_conv2d_nchw, has_groups=True),
205+
wrap_topi_schedule(topi.cuda.schedule_group_conv2d_nchw),
206+
name="group_conv2d_nchw.cuda",
207+
)
208+
else:
209+
raise RuntimeError(
210+
"General group convolution is not currently supported for NCHWc layouts"
211+
)
185212
else:
186-
raise RuntimeError("General group convolution is not currently supported")
213+
raise RuntimeError(
214+
"General group convolution has limited support for NCHW(4c) layouts..."
215+
)
187216
return strategy
188217

189218

python/tvm/topi/adreno/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
from .conv2d_nchw import *
2121
from .depthwise_conv2d_nchw import *
2222
from .conv2d_nhwc import *
23+
from .group_conv2d_nchw import *
2324
from .depthwise_conv2d_nhwc import *
2425
from .pooling import *
2526
from .conv2d_alter_op import *

0 commit comments

Comments
 (0)