jverma-quic commented on code in PR #11417:
URL: https://github.com/apache/tvm/pull/11417#discussion_r883864069
##########
tests/python/contrib/test_hexagon/infrastructure.py:
##########
@@ -228,3 +227,18 @@ def compute(n, ho, wo, ko, hi, wi, ki):
)
return output_shape, compute
+
+
+# Transpose and reshape numpy array according to the specified layout
+def transform_numpy(arr_np, layout):
Review 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)`?
You're right that this function is making an assumption about the supplied
layout which can transform the input incorrectly. I will include the
current_layout as an argument. Thanks for the suggestion.
##########
python/tvm/topi/hexagon/slice_ops/avg_pool2d.py:
##########
@@ -0,0 +1,198 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from tvm.ir.module import IRModule
+from tvm import te
+from tvm import tir
+from tvm.script import tir as T
+from ..utils import apply_transform, get_layout_transform_fn
+
+
+# The slice op implementation for avg_pool2d makes serveral assumptions:
+# 1) Both input and output are a multiple of croutons, and the input is already
+# padded for a given output shape as per any crouton and non-crouton related
+# padding.
+# 2) The current implementation assumes 'count_include_pad' to be 'True'. It
can
+# modified to support 'False' but the element count for the pooling window
must
+# be pre-computed and provided as an input to reduce the run-time overhead.
+# 3) 'padding' is also ignored. It must be handled outside of the sliced op.
+# 4) Please note that this implementation will not work if the output was
padded
+# for the croutons. Since we loop over the logical output shape, this can
result
+# into out-of-bound access for the input.
+
+def avg_pool2d_compute(A, out_shape, kernel, stride, dilation):
+ kh, kw = kernel
+ rh = te.reduce_axis((0, kh), name="rh")
+ rw = te.reduce_axis((0, kw), name="rw")
+ ob, oh, ow, oc = out_shape
+ sh, sw = stride
+ dh, dw = dilation
+ Area = float(1) / (kh * kw)
+
+ Sum = te.compute(
+ out_shape,
+ lambda b, h, w, c: te.sum(
+ A[b, h * sh + dh * rh, w * sw + dw * rw, c].astype("float32"),
axis=[rh, rw]
+ ),
+ name="sum",
+ )
+ Avg = te.compute(
+ out_shape, lambda b, h, w, c: (Sum[b, h, w, c] *
Area).astype(A.dtype), name="avg"
+ )
+ return Avg
+
+
+# Schedule for input and output layout nhwc-8h2w32c2w
+def STIR_schedule_nhwc_8h2w32c2w(outs, ins, output_layout: str, input_layout:
str):
+ func = te.create_prim_func([ins, outs])
+ s = tir.Schedule(func)
+ Sum = s.get_block("sum")
+ Avg = s.get_block("avg")
+
+ apply_transform(s, Sum, 0, "read", input_layout)
+ apply_transform(s, Avg, 0, "write", output_layout)
+
+ # Schedule 'Sum'
+ bn, bh, bw, bc, rx, ry = s.get_loops(Sum)
+ bho, bhi = s.split(bh, [None, 8])
+ 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
Review 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.
--
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-unsubscr...@tvm.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
