Page MenuHomePhabricator

[mlir][tosa] Migrate tosa to more efficient linalg.conv
ClosedPublic

Authored by rsuderman on Aug 4 2021, 3:03 PM.

Details

Summary

Existing linalg.conv2d is not well optimized for performance. Changed to a
version that is more aligned for optimziation. Include the corresponding
transposes to use this optimized version.

This also splits the conv and depthwise conv into separate implementations
to avoid overly complex lowerings.

Diff Detail

Event Timeline

rsuderman created this revision.Aug 4 2021, 3:03 PM
rsuderman requested review of this revision.Aug 4 2021, 3:03 PM
rsuderman updated this revision to Diff 364260.Aug 4 2021, 3:04 PM

Removed no longer needed convolutions.

rsuderman updated this revision to Diff 364262.Aug 4 2021, 3:06 PM

Removed no longer needed named-ops.mlir test case.

rsuderman updated this revision to Diff 364279.Aug 4 2021, 4:02 PM

Fixed transpose conv test.

antiagainst requested changes to this revision.Aug 5 2021, 6:25 AM
antiagainst added inline comments.
mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
845

This op is meant to replace the existing conv_2d_input_nhwc_filter_hwcf right? We'd better to follow exactly how the loops are mapped so that downstream users of conv_2d_input_nhwc_filter_hwcf won't be surprised to see loop order change (which has implications on performance) after migration.

In conv_2d_input_nhwc_filter_hwcf, the first four loops corresponds to N, OH, OW, OC, F (https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsSpec.tc#L266), but here it's N, F, H, W?

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
993

Would be nice to have a constant folder for the case where filter is constant (which is common for inference.)

1087–1088

NHWF is more natural. I think we can define conv_2d_nhwc_hwcf to have NHWF output format (which follows the existing conv_2d_input_nhwc_filter_hwcf) and to remove this transpose?

This revision now requires changes to proceed.Aug 5 2021, 6:25 AM
rsuderman updated this revision to Diff 364611.Aug 5 2021, 2:15 PM

Swapped F to final channel of output.

rsuderman updated this revision to Diff 364612.Aug 5 2021, 2:17 PM

Added comment about folding constant.

rsuderman marked 3 inline comments as done.Aug 5 2021, 2:19 PM
rsuderman added inline comments.
mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
845

Oops, I think I looked at the wrong conv2D when I was porting this. Two adjacent ones are named very similarly. (Thought it was weird but didn't question).

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
993

Added a comment. If I add a tosa::Transpose constant folder it should be easy to integrate however I need to verify this constant is not used elsewhere. Otherwise we risk doubling up the filter constant in model training.

rsuderman updated this revision to Diff 364635.Aug 5 2021, 3:26 PM
rsuderman marked 2 inline comments as done.

Fixed output dimension ordering.

rsuderman updated this revision to Diff 364639.Aug 5 2021, 3:31 PM

remove unused

antiagainst requested changes to this revision.Aug 10 2021, 2:42 PM
antiagainst added inline comments.
mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
631–632

This needs to be regenerated?

mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py
193

This should have the first 4 dims as parallel too? Basically the first 4 dims should be the 4 dims from output sequentially.

See here for the full order: https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsSpec.tc#L212

This revision now requires changes to proceed.Aug 10 2021, 2:42 PM
rsuderman updated this revision to Diff 365623.Aug 10 2021, 3:56 PM

synced to head and fixed some naming issues

rsuderman marked 2 inline comments as done.Aug 10 2021, 4:19 PM
rsuderman added inline comments.
mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py
193

Its hard to tell if this is specifically related to depthwise conv. For now I am focusing on the convolution refactor (and can fix depthwise in another CL). With the regenerated code its much easier to address these kinds of changes in the source python.

antiagainst accepted this revision.Aug 11 2021, 6:32 AM

LGTM! Just a few nits.

mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py
148–151

Might be good to name this as conv_2d_nchw_fchw for consistency?

174

It would be nice to list the implicit loop order in the doc like what we do for the TC ops.

This revision is now accepted and ready to land.Aug 11 2021, 6:32 AM
This revision was automatically updated to reflect the committed changes.
rsuderman marked an inline comment as done.