Currently there is no way to initialize the output tensor of a reduction with
another Tensor.
One use case for such a functionality is when we have a fused **convolution
with bias-add** operation. Currently define 2 compute operations, where the
first one defines the reduction in the convolution, and another compute for
bias add as shown below:
Output = te.compute(
(batch, out_height, out_width, out_channel),
lambda nn, yy, xx, ff: te.sum(
PaddedInput[nn, yy * stride_h + ry * dilation_h,
xx * stride_w + rx * dilation_w, rc].astype(out_dtype) *
Filter[ry, rx, rc, ff].astype(out_dtype), axis=[ry, rx, rc]),
name="Conv2dOutput", tag="conv2d_nhwc")
Output_with_bias = te.compute(
(batch, out_height, out_width, out_channel),
lambda nn, yy, xx, ff: Output[nn, yy, xx, ff] + Bias[nn, yy, xx, ff])
The code generated for the above computation would contain a fused loop nest,
which has an initialization with `0f` operation for the `Output` tensor,
followed by the actual reduction for the convolution and then adding bias-add
with Output. What I wanted to do is to remove the bias add computation and
instead initialize the `Output` Tensor with the `Bias` values.
One way to do this would be to add an `init` member to `ReduceNode` object and
an `init` argument to the `reducer` function returned by `CommReducer`, which
can be used to pass in a Tensor with which to initialize the output Tensor
instead of the identity element.
For Example, `te.sum` would be used as shown below:
te.sum(
PaddedInput[nn, yy * stride_h + ry * dilation_h,
xx * stride_w + rx * dilation_w, rc].astype(out_dtype) *
Filter[ry, rx, rc, ff].astype(out_dtype), axis=[ry, rx, rc],
init=Bias[nn, yy, xx, ff])
`init=Bias[nn, yy, xx, ff]` passes the Tensor that can be used in
`MakeReduction` function in `compute_op.cc` to initialize with `Bias` instead
of the identity element.
Let me know what you think about adding this option and any suggestions about
my approach to do this.
Thanks,
Anirudh
---
[Visit
Topic](https://discuss.tvm.ai/t/add-init-option-to-reducenode-to-initialize-with-custom-tensors/6944/1)
to respond.
You are receiving this because you enabled mailing list mode.
To unsubscribe from these emails, [click
here](https://discuss.tvm.ai/email/unsubscribe/510ec2c24eabfaec3de8ca819aa1b64288a3d0bd84888e1ad7b1e0781e1738ea).
