Add `init` member to ReduceNode (#6138)
authorquic-sanirudh <63797228+quic-sanirudh@users.noreply.github.com>
Thu, 27 Aug 2020 02:11:24 +0000 (07:41 +0530)
committerGitHub <noreply@github.com>
Thu, 27 Aug 2020 02:11:24 +0000 (19:11 -0700)
commitc6dd26b9d73c60b3122bafe800529661e14f75e1
tree3893ca56db6e16def3f73190bb0c3dcad8eff4fd
parent415c0883ab101be9854315f84ae9e1e5e4295563
Add `init` member to ReduceNode (#6138)

- This patch adds a new member to ReduceNode called init which allows
  initialization with a custom ProducerLoad or a Float/Int immediate.
- This allows initialization of the output Tensor of a reduction with
  another Tensor instead of the `identity_element` defined in the
  CommReducer
- One example use case for this node is to initialize the Output of a
  convolution reduction with the Bias values thereby saving the
  Bias-add computation.
20 files changed:
include/tvm/tir/expr.h
include/tvm/tir/op.h
include/tvm/topi/reduction.h
python/tvm/tir/expr.py
python/tvm/tir/op.py
src/arith/canonical_simplify.cc
src/printer/tir_text_printer.cc
src/te/autodiff/ad_simplify.cc
src/te/autodiff/ad_util.cc
src/te/autodiff/jacobian.cc
src/te/operation/compute_op.cc
src/te/operation/cross_thread_reduction.cc
src/te/operation/tensorize.cc
src/te/schedule/schedule_dataflow_rewrite.cc
src/tir/ir/expr.cc
src/tir/ir/expr_functor.cc
src/tir/op/op.cc
tests/python/integration/test_reduce.py
tests/python/unittest/test_arith_canonical_simplify.py
tests/python/unittest/test_te_autodiff.py