You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by "lhutton1 (via GitHub)" <gi...@apache.org> on 2023/06/19 09:42:08 UTC
[GitHub] [tvm] lhutton1 commented on a diff in pull request #14855: [topi] Add `arm_cpu` specific pooling schedules
lhutton1 commented on code in PR #14855:
URL: https://github.com/apache/tvm/pull/14855#discussion_r1233751645
##########
python/tvm/relay/op/strategy/arm_cpu.py:
##########
@@ -63,19 +64,12 @@ def concatenate_strategy_arm_cpu(attrs, inputs, out_type, target):
def schedule_pool_arm_cpu(attrs, outs, target):
"""schedule pooling ops arm cpu"""
layout = attrs.layout
- avg_pool = isinstance(attrs, relay.op.op_attrs.AvgPool2DAttrs)
with target:
- if (
- avg_pool
- and target.features.has_dsp
- and layout in ("NCW", "NCHW")
- or not avg_pool
- and target.features.has_dsp
- and layout in ("NWC", "NHWC")
- ):
+ if target.features.has_dsp:
+ is_avg_pool = isinstance(attrs, relay.op.op_attrs.AvgPool2DAttrs)
+ return dsp.pool.schedule_pool(outs, layout, is_avg_pool)
+ else:
Review Comment:
nit: `else` can be removed
##########
python/tvm/topi/arm_cpu/pooling.py:
##########
@@ -17,9 +17,88 @@
# pylint: disable=invalid-name, unused-variable
"""Schedule for pooling operators"""
-from .mprofile.dsp.pool import pool_dsp_schedule
+import logging
+from tvm import topi, te
+from tvm.target import Target
+from .. import tag
def schedule_pool(outs, layout):
- """Create schedule for avgpool/maxpool with dsp"""
- return pool_dsp_schedule(outs, layout)
+ """Create schedule for avgpool/maxpool"""
+
+ if layout != "NHWC":
+ logger = logging.getLogger("topi")
+ logger.warning(
+ """We currently only support NHWC target specific pools on arm_cpu,
+ falling back on generic pool scheduling"""
+ )
+ return topi.generic.schedule_pool(outs, layout)
+
+ return schedule_pool_2d(outs)
+
+
+def schedule_pool_2d(outs):
+ """Create arm_cpu specific 2D schedule for avgpool/maxpool"""
+
+ outs = [outs] if isinstance(outs, te.tensor.Tensor) else outs
+ schedule_ops = [x.op for x in outs]
+ schedule = te.create_schedule(schedule_ops)
+ scheduled_ops = []
+
+ def traverse(op):
+ # Recursively inline any injective operation that isn't the pooling
+ # operation or hasn't already been scheduled.
+ if tag.is_injective(op.tag):
+ if op not in schedule.outputs:
+ schedule[op].compute_inline()
+ for tensor in op.input_tensors:
+ if isinstance(tensor.op, te.tensor.ComputeOp) and tensor.op not in scheduled_ops:
+ traverse(tensor.op)
+ # schedule the actual pooling operation
+ elif op.tag.startswith("pool"):
+ n, height, width, channel = schedule[op].op.axis
+ # Average pool consists of two parts; a sum then a division.
+ # We can schedule the division loop to parallelize across height and
+ # vectorize across width.
+ if op != outs[0].op:
+ output = outs[0]
+ output_fused = schedule[output].fuse(output.op.axis[1], output.op.axis[2])
+ schedule[output].parallel(output_fused)
+ _, inner = schedule[output].split(output.op.axis[3], 8)
+ schedule[output].vectorize(inner)
+
+ enable_explicit_vectorization = not Target.current(allow_none=False).features.has_sve
+ PaddedInput = op.input_tensors[0]
+ if isinstance(PaddedInput.op, te.tensor.ComputeOp):
+ schedule[PaddedInput].compute_inline()
+
+ # For targets without SVE try explicitly vectorizing the channel
+ # loop, For SVE targets leave the loop in place for LLVM to convert
+ # into a scalable vector loop.
+ vectorization_factor = 8 if enable_explicit_vectorization else channel.dom.extent
Review Comment:
Could this factor be tuned with AutoTVM at some point?
##########
python/tvm/topi/arm_cpu/pooling.py:
##########
@@ -17,9 +17,88 @@
# pylint: disable=invalid-name, unused-variable
"""Schedule for pooling operators"""
-from .mprofile.dsp.pool import pool_dsp_schedule
+import logging
+from tvm import topi, te
+from tvm.target import Target
+from .. import tag
def schedule_pool(outs, layout):
- """Create schedule for avgpool/maxpool with dsp"""
- return pool_dsp_schedule(outs, layout)
+ """Create schedule for avgpool/maxpool"""
+
+ if layout != "NHWC":
+ logger = logging.getLogger("topi")
+ logger.warning(
+ """We currently only support NHWC target specific pools on arm_cpu,
+ falling back on generic pool scheduling"""
+ )
+ return topi.generic.schedule_pool(outs, layout)
+
+ return schedule_pool_2d(outs)
+
+
+def schedule_pool_2d(outs):
+ """Create arm_cpu specific 2D schedule for avgpool/maxpool"""
+
+ outs = [outs] if isinstance(outs, te.tensor.Tensor) else outs
+ schedule_ops = [x.op for x in outs]
+ schedule = te.create_schedule(schedule_ops)
+ scheduled_ops = []
+
+ def traverse(op):
+ # Recursively inline any injective operation that isn't the pooling
+ # operation or hasn't already been scheduled.
+ if tag.is_injective(op.tag):
+ if op not in schedule.outputs:
+ schedule[op].compute_inline()
+ for tensor in op.input_tensors:
+ if isinstance(tensor.op, te.tensor.ComputeOp) and tensor.op not in scheduled_ops:
+ traverse(tensor.op)
+ # schedule the actual pooling operation
+ elif op.tag.startswith("pool"):
+ n, height, width, channel = schedule[op].op.axis
+ # Average pool consists of two parts; a sum then a division.
+ # We can schedule the division loop to parallelize across height and
+ # vectorize across width.
+ if op != outs[0].op:
+ output = outs[0]
+ output_fused = schedule[output].fuse(output.op.axis[1], output.op.axis[2])
+ schedule[output].parallel(output_fused)
+ _, inner = schedule[output].split(output.op.axis[3], 8)
+ schedule[output].vectorize(inner)
Review Comment:
Could SVE be used here?
##########
python/tvm/topi/arm_cpu/pooling.py:
##########
@@ -17,9 +17,88 @@
# pylint: disable=invalid-name, unused-variable
"""Schedule for pooling operators"""
-from .mprofile.dsp.pool import pool_dsp_schedule
+import logging
+from tvm import topi, te
+from tvm.target import Target
+from .. import tag
def schedule_pool(outs, layout):
- """Create schedule for avgpool/maxpool with dsp"""
- return pool_dsp_schedule(outs, layout)
+ """Create schedule for avgpool/maxpool"""
+
+ if layout != "NHWC":
+ logger = logging.getLogger("topi")
+ logger.warning(
+ """We currently only support NHWC target specific pools on arm_cpu,
+ falling back on generic pool scheduling"""
+ )
+ return topi.generic.schedule_pool(outs, layout)
+
+ return schedule_pool_2d(outs)
+
+
+def schedule_pool_2d(outs):
+ """Create arm_cpu specific 2D schedule for avgpool/maxpool"""
+
+ outs = [outs] if isinstance(outs, te.tensor.Tensor) else outs
+ schedule_ops = [x.op for x in outs]
+ schedule = te.create_schedule(schedule_ops)
+ scheduled_ops = []
+
+ def traverse(op):
+ # Recursively inline any injective operation that isn't the pooling
+ # operation or hasn't already been scheduled.
+ if tag.is_injective(op.tag):
+ if op not in schedule.outputs:
+ schedule[op].compute_inline()
+ for tensor in op.input_tensors:
+ if isinstance(tensor.op, te.tensor.ComputeOp) and tensor.op not in scheduled_ops:
+ traverse(tensor.op)
+ # schedule the actual pooling operation
+ elif op.tag.startswith("pool"):
+ n, height, width, channel = schedule[op].op.axis
+ # Average pool consists of two parts; a sum then a division.
+ # We can schedule the division loop to parallelize across height and
+ # vectorize across width.
+ if op != outs[0].op:
+ output = outs[0]
+ output_fused = schedule[output].fuse(output.op.axis[1], output.op.axis[2])
+ schedule[output].parallel(output_fused)
+ _, inner = schedule[output].split(output.op.axis[3], 8)
+ schedule[output].vectorize(inner)
+
+ enable_explicit_vectorization = not Target.current(allow_none=False).features.has_sve
+ PaddedInput = op.input_tensors[0]
+ if isinstance(PaddedInput.op, te.tensor.ComputeOp):
+ schedule[PaddedInput].compute_inline()
+
+ # For targets without SVE try explicitly vectorizing the channel
+ # loop, For SVE targets leave the loop in place for LLVM to convert
+ # into a scalable vector loop.
+ vectorization_factor = 8 if enable_explicit_vectorization else channel.dom.extent
+ channel_outer, channel_inner = schedule[op].split(channel, vectorization_factor)
+ schedule[op].vectorize(channel_inner)
+ schedule[op].parallel(height)
+ if len(schedule[op].op.reduce_axis) > 0:
+ filter_height, filter_width = schedule[op].op.reduce_axis
+ if filter_height.dom.extent * filter_width.dom.extent <= 9:
+ # For small filters, unrolling the filter loops allows us to
+ # vectorize over channels without reordering anything.
+ schedule[op].unroll(filter_width)
+ schedule[op].unroll(filter_height)
+ else:
Review Comment:
Do larger filter sizes get tested at all?
##########
python/tvm/topi/arm_cpu/pooling.py:
##########
@@ -17,9 +17,88 @@
# pylint: disable=invalid-name, unused-variable
"""Schedule for pooling operators"""
-from .mprofile.dsp.pool import pool_dsp_schedule
+import logging
+from tvm import topi, te
+from tvm.target import Target
+from .. import tag
def schedule_pool(outs, layout):
- """Create schedule for avgpool/maxpool with dsp"""
- return pool_dsp_schedule(outs, layout)
+ """Create schedule for avgpool/maxpool"""
+
+ if layout != "NHWC":
+ logger = logging.getLogger("topi")
+ logger.warning(
+ """We currently only support NHWC target specific pools on arm_cpu,
+ falling back on generic pool scheduling"""
+ )
+ return topi.generic.schedule_pool(outs, layout)
+
+ return schedule_pool_2d(outs)
+
+
+def schedule_pool_2d(outs):
+ """Create arm_cpu specific 2D schedule for avgpool/maxpool"""
+
+ outs = [outs] if isinstance(outs, te.tensor.Tensor) else outs
+ schedule_ops = [x.op for x in outs]
+ schedule = te.create_schedule(schedule_ops)
+ scheduled_ops = []
+
+ def traverse(op):
+ # Recursively inline any injective operation that isn't the pooling
+ # operation or hasn't already been scheduled.
+ if tag.is_injective(op.tag):
+ if op not in schedule.outputs:
+ schedule[op].compute_inline()
+ for tensor in op.input_tensors:
+ if isinstance(tensor.op, te.tensor.ComputeOp) and tensor.op not in scheduled_ops:
+ traverse(tensor.op)
+ # schedule the actual pooling operation
+ elif op.tag.startswith("pool"):
+ n, height, width, channel = schedule[op].op.axis
+ # Average pool consists of two parts; a sum then a division.
+ # We can schedule the division loop to parallelize across height and
+ # vectorize across width.
+ if op != outs[0].op:
+ output = outs[0]
+ output_fused = schedule[output].fuse(output.op.axis[1], output.op.axis[2])
+ schedule[output].parallel(output_fused)
+ _, inner = schedule[output].split(output.op.axis[3], 8)
+ schedule[output].vectorize(inner)
+
+ enable_explicit_vectorization = not Target.current(allow_none=False).features.has_sve
+ PaddedInput = op.input_tensors[0]
Review Comment:
nit: better to use lower case for variable names e.g. `padded_input`
##########
python/tvm/topi/arm_cpu/pooling.py:
##########
@@ -17,9 +17,88 @@
# pylint: disable=invalid-name, unused-variable
"""Schedule for pooling operators"""
-from .mprofile.dsp.pool import pool_dsp_schedule
+import logging
+from tvm import topi, te
+from tvm.target import Target
+from .. import tag
def schedule_pool(outs, layout):
- """Create schedule for avgpool/maxpool with dsp"""
- return pool_dsp_schedule(outs, layout)
+ """Create schedule for avgpool/maxpool"""
+
+ if layout != "NHWC":
+ logger = logging.getLogger("topi")
+ logger.warning(
+ """We currently only support NHWC target specific pools on arm_cpu,
+ falling back on generic pool scheduling"""
+ )
+ return topi.generic.schedule_pool(outs, layout)
+
+ return schedule_pool_2d(outs)
+
+
+def schedule_pool_2d(outs):
+ """Create arm_cpu specific 2D schedule for avgpool/maxpool"""
+
+ outs = [outs] if isinstance(outs, te.tensor.Tensor) else outs
+ schedule_ops = [x.op for x in outs]
+ schedule = te.create_schedule(schedule_ops)
+ scheduled_ops = []
+
+ def traverse(op):
+ # Recursively inline any injective operation that isn't the pooling
+ # operation or hasn't already been scheduled.
+ if tag.is_injective(op.tag):
+ if op not in schedule.outputs:
+ schedule[op].compute_inline()
+ for tensor in op.input_tensors:
+ if isinstance(tensor.op, te.tensor.ComputeOp) and tensor.op not in scheduled_ops:
+ traverse(tensor.op)
+ # schedule the actual pooling operation
+ elif op.tag.startswith("pool"):
+ n, height, width, channel = schedule[op].op.axis
+ # Average pool consists of two parts; a sum then a division.
+ # We can schedule the division loop to parallelize across height and
+ # vectorize across width.
+ if op != outs[0].op:
+ output = outs[0]
+ output_fused = schedule[output].fuse(output.op.axis[1], output.op.axis[2])
+ schedule[output].parallel(output_fused)
+ _, inner = schedule[output].split(output.op.axis[3], 8)
+ schedule[output].vectorize(inner)
+
+ enable_explicit_vectorization = not Target.current(allow_none=False).features.has_sve
+ PaddedInput = op.input_tensors[0]
+ if isinstance(PaddedInput.op, te.tensor.ComputeOp):
+ schedule[PaddedInput].compute_inline()
+
+ # For targets without SVE try explicitly vectorizing the channel
+ # loop, For SVE targets leave the loop in place for LLVM to convert
+ # into a scalable vector loop.
+ vectorization_factor = 8 if enable_explicit_vectorization else channel.dom.extent
+ channel_outer, channel_inner = schedule[op].split(channel, vectorization_factor)
+ schedule[op].vectorize(channel_inner)
+ schedule[op].parallel(height)
+ if len(schedule[op].op.reduce_axis) > 0:
+ filter_height, filter_width = schedule[op].op.reduce_axis
+ if filter_height.dom.extent * filter_width.dom.extent <= 9:
+ # For small filters, unrolling the filter loops allows us to
+ # vectorize over channels without reordering anything.
+ schedule[op].unroll(filter_width)
+ schedule[op].unroll(filter_height)
+ else:
+ # Reordering so that channels is the fastest moving axis allows
+ # LLVM to vectorize across contiguous memory in the NHWC
+ # ordering.
+ schedule[op].reorder(
+ n, height, width, filter_height, filter_width, channel_outer, channel_inner
+ )
+ else:
+ schedule[op].reorder(n, height, width, channel_outer, channel_inner)
Review Comment:
Is this reachable?
##########
python/tvm/topi/arm_cpu/pooling.py:
##########
@@ -17,9 +17,88 @@
# pylint: disable=invalid-name, unused-variable
"""Schedule for pooling operators"""
-from .mprofile.dsp.pool import pool_dsp_schedule
+import logging
+from tvm import topi, te
+from tvm.target import Target
+from .. import tag
def schedule_pool(outs, layout):
- """Create schedule for avgpool/maxpool with dsp"""
- return pool_dsp_schedule(outs, layout)
+ """Create schedule for avgpool/maxpool"""
+
+ if layout != "NHWC":
+ logger = logging.getLogger("topi")
+ logger.warning(
+ """We currently only support NHWC target specific pools on arm_cpu,
+ falling back on generic pool scheduling"""
+ )
+ return topi.generic.schedule_pool(outs, layout)
+
+ return schedule_pool_2d(outs)
+
+
+def schedule_pool_2d(outs):
+ """Create arm_cpu specific 2D schedule for avgpool/maxpool"""
+
+ outs = [outs] if isinstance(outs, te.tensor.Tensor) else outs
+ schedule_ops = [x.op for x in outs]
+ schedule = te.create_schedule(schedule_ops)
+ scheduled_ops = []
+
+ def traverse(op):
+ # Recursively inline any injective operation that isn't the pooling
+ # operation or hasn't already been scheduled.
+ if tag.is_injective(op.tag):
+ if op not in schedule.outputs:
+ schedule[op].compute_inline()
+ for tensor in op.input_tensors:
+ if isinstance(tensor.op, te.tensor.ComputeOp) and tensor.op not in scheduled_ops:
+ traverse(tensor.op)
+ # schedule the actual pooling operation
+ elif op.tag.startswith("pool"):
+ n, height, width, channel = schedule[op].op.axis
+ # Average pool consists of two parts; a sum then a division.
+ # We can schedule the division loop to parallelize across height and
+ # vectorize across width.
+ if op != outs[0].op:
+ output = outs[0]
+ output_fused = schedule[output].fuse(output.op.axis[1], output.op.axis[2])
+ schedule[output].parallel(output_fused)
+ _, inner = schedule[output].split(output.op.axis[3], 8)
+ schedule[output].vectorize(inner)
+
+ enable_explicit_vectorization = not Target.current(allow_none=False).features.has_sve
+ PaddedInput = op.input_tensors[0]
+ if isinstance(PaddedInput.op, te.tensor.ComputeOp):
+ schedule[PaddedInput].compute_inline()
+
+ # For targets without SVE try explicitly vectorizing the channel
+ # loop, For SVE targets leave the loop in place for LLVM to convert
+ # into a scalable vector loop.
+ vectorization_factor = 8 if enable_explicit_vectorization else channel.dom.extent
+ channel_outer, channel_inner = schedule[op].split(channel, vectorization_factor)
+ schedule[op].vectorize(channel_inner)
+ schedule[op].parallel(height)
+ if len(schedule[op].op.reduce_axis) > 0:
+ filter_height, filter_width = schedule[op].op.reduce_axis
+ if filter_height.dom.extent * filter_width.dom.extent <= 9:
Review Comment:
Curious if there is any reason for "9" being chosen as the cut-off
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org