Add CostInferenceFunction for SplitOp (#63133)
authorTanvir Zaman <motanv@fb.com>
Fri, 13 Aug 2021 19:25:16 +0000 (12:25 -0700)
committerFacebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
Fri, 13 Aug 2021 19:28:15 +0000 (12:28 -0700)
commitcc6b023cbae043ae30bdb1da743d5aec5afce02b
treebc165978e8eb3f1d9485287adf094a06b06dfdc7
parentacdad8bc6352c4f35bf4974c2dce8b8c8019841d
Add CostInferenceFunction for SplitOp (#63133)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/63133

SplitOp is costly but missing cost inference function which hurts cost based balancing. Changes are:
(1) Addition of CostInferenceFunction for SplitOp
(2) Small fix in CostInferenceFunction for ConcatOp

Test Plan:
Added unit tests:

buck test //caffe2/caffe2/python/operator_test:split_op_cost_test

buck test //caffe2/caffe2/python/operator_test:concat_op_cost_test

Reviewed By: smacke

Differential Revision: D30247360

fbshipit-source-id: 989e962f3a981acc85b73aac3fb23e603b7d1591
caffe2/operators/concat_split_op.cc
caffe2/operators/concat_split_op.h
caffe2/python/operator_test/concat_op_cost_test.py [new file with mode: 0644]
caffe2/python/operator_test/split_op_cost_test.py [new file with mode: 0644]