[Hexagon] Implement avg_pool2d slice op #11417
Conversation
|
@csullivan @Lunderberg @cconvey, please review these changes. |
python/tvm/topi/hexagon/utils.py
Outdated
| def n11c_1024c_2d(n, h, w, c): | ||
| return [n, h, w, c // 1024, te.AXIS_SEPARATOR, c % 1024] | ||
|
|
||
|
|
||
| def n11c_1024c_1d(n, h, w, c): | ||
| return [n, h, w, c // 1024, c % 1024] | ||
|
|
||
|
|
||
| def nhwc_8h2w32c2w_2d(n, h, w, c): | ||
| return [n, h // 8, w // 4, c // 32, te.AXIS_SEPARATOR, h % 8, (w % 4) // 2, c % 32, w % 2] | ||
|
|
||
|
|
||
| def nhwc_8h2w32c2w_1d(n, h, w, c): | ||
| return [n, h // 8, w // 4, c // 32, h % 8, (w % 4) // 2, c % 32, w % 2] |
There was a problem hiding this comment.
Would these functions' purpose be clearer if their names started with something like get_shape_..., xform_layout_..., etc?
There was a problem hiding this comment.
Maybe. Although, I would prefer the current names.
|
|
||
|
|
||
| # Transpose and reshape numpy array according to the specified layout | ||
| def transform_numpy(arr_np, layout): |
There was a problem hiding this comment.
The function seems to assume that the supplied layout is NHWC. Is that a safe assumption for all expected uses of the function?
If no, then should we put nhwc into the function name, or perhaps change its argument list to something like (arr_np, current_layout, new_layout)?
There was a problem hiding this comment.
You're right that this function is making an assumption about the supplied layout which can transform the input incorrectly. Thanks for the suggestion, @cconvey!
| if output_layout == "n11c-1024c": | ||
| assert ( | ||
| pad_before_w == 0 and pad_after_w == 0 and pad_before_h == 0 and pad_after_h == 0 | ||
| ), "Padding must be zero for n11c-1024c layout!!" |
There was a problem hiding this comment.
Nit: (at most) one exclamation point is probably enough here. Having this appear in an assertion-failed message is probably enough to get the user's attention.
python/tvm/topi/hexagon/utils.py
Outdated
| else: | ||
| raise RuntimeError(f"Unexpected layout '{layout}'") | ||
|
|
||
| s.transform_layout(block, block_index, buffer_type, transform_fn) |
There was a problem hiding this comment.
FYI, after #11269 lands, the calling layout_transform will also handle the call to set_axis_separators, so this function may become simpler or empty.
| from ..utils import apply_transform, get_layout_transform_fn | ||
|
|
||
|
|
||
| # The slice op implementation for avg_pool2d makes serveral assumptions: |
There was a problem hiding this comment.
Can this be moved from a comment to a docstring?
| name="sum", | ||
| ) | ||
| Avg = te.compute( | ||
| out_shape, lambda b, h, w, c: (Sum[b, h, w, c] * Area).astype(A.dtype), name="avg" |
There was a problem hiding this comment.
Nitpick: The name Area threw me a bit, as I initially thought Area should be the area of the kernel relative to a single value, rather than the area of a value relative to the kernel. Can we rename Area to either InvArea or NumValues?
| bwo, bwi = s.split(bw, [None, 4]) | ||
| bwio, bwii = s.split(bwi, [None, 2]) # Doesn't seem to be doing anything | ||
| bco, bci = s.split(bc, [None, 32]) | ||
| s.reorder(bn, bho, bwo, bco, bhi, bwio, rx, ry, bci, bwii) # --- DOESN'T do anything |
There was a problem hiding this comment.
What do you have before and after these lines? Running the test case test_avg_pool2d_slice.py::TestAvgPool2dSlice::test_avg_pool2d_slice[nhwc-8h2w32c2w-False-str ide0-kernel0-float16-dilation0-padding0-True-nhwc-8h2w32c2w-output_shape0-False] and using print(s.mod.script()), I can see the loopnest before this line to have extents T.grid(1, 1, 8, 2, 2, 2, 1, 32, 3, 3) and afterward to have extents T.grid(1, 1, 2, 1, 8, 2, 3, 3, 32, 2), so it does look like the reorder is having an effect.
There was a problem hiding this comment.
You're right. I do see the loops getting reordered after this line. However, when I print it again after s.compute_at(Sum, hi), I don't see the reordered/fused loopnest anymore.
| bwio, bwii = s.split(bwi, [None, 2]) # Doesn't seem to be doing anything | ||
| bco, bci = s.split(bc, [None, 32]) | ||
| s.reorder(bn, bho, bwo, bco, bhi, bwio, rx, ry, bci, bwii) # --- DOESN'T do anything | ||
| bci_wii = s.fuse(bci, bwii) # --- DOESN'T do anything |
There was a problem hiding this comment.
Same question here, after fusing I see extents T.grid(1, 1, 2, 1, 8, 2, 3, 3, 64) and can't reproduce the lack of effect.
Lunderberg
changed the title
| hexagon_session, | ||
| ): | ||
|
|
||
| target_hexagon = tvm.target.hexagon("v69") |
There was a problem hiding this comment.
AFAIK, TVM's Hexagon CI has only tested v68 code in the past. Is there any possibility that specifying v69 here will break CI or require some newer version of the Hexagon SDK?
There was a problem hiding this comment.
I think it should be okay but will confirm at our end. BTW, do you know which Hexagon SDK is used for upstream CI?
|
Hi @Lunderberg and @cconvey, please let me know if there are any additional comments. |
| @@ -0,0 +1,369 @@ | |||
| # Licensed to the Apache Software Foundation (ASF) under one | |||
There was a problem hiding this comment.
please move this file to tests/python/contrib/test_hexagon/topi
| s.transform_layout(Sum, ("read", 0), input_transform_fn) | ||
| s.transform_layout(Avg, ("write", 0), output_transform_fn) |
There was a problem hiding this comment.
Not necessary for this PR but just sharing: there's a new API sugar for transform_layout that allows you to address the block and buffer by name, e.g.
sch.transform_layout(block="tensor", buffer="placeholder", index_map=_int8_nhwc_8h8w32c_map)
There was a problem hiding this comment.
Thanks for the suggestion, @csullivan! I have tried it in the past but couldn't get the new API sugar to work due to the intermediate compute in avgpool.
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| sys.exit(pytest.main(sys.argv)) |
There was a problem hiding this comment.
Don't need to update the PR only for this, but if you're adding more commit please change this line to:
tvm.testing.main()
Thanks for contributing to TVM! Please refer to guideline https://tvm.apache.org/docs/contribute/ for useful information and tips. After the pull request is submitted, please request code reviews from Reviewers by @ them in the pull request thread.
cc @mehrdadh