Fix bytes_written and bytes_read (#64040)
authorTanvir Zaman <motanv@fb.com>
Mon, 30 Aug 2021 19:56:15 +0000 (12:56 -0700)
committerFacebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
Mon, 30 Aug 2021 19:57:31 +0000 (12:57 -0700)
commite98173ff3423247c597e21c923c8f47470ef07ab
tree2613870e901e79905e2489076a58768ca81d8ad2
parenteafe33c995d47d45fceaf42801717f3120d799b9
Fix bytes_written and bytes_read (#64040)

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

In operator cost inference functions, in many places we are using sizeof(x.data_type()). Since data_type() returns a 32 bit integer from [this enum](https://www.internalfb.com/code/fbsource/[15e7ffe4073cf08c61077c7c24a4839504b964a2]/fbcode/caffe2/caffe2/proto/caffe2.proto?lines=20), we are basically always getting 4 for sizeof(x.data_type()) no matter what actual data type x has. Big thanks to Jack Langman for specifically pointing to this bug.

We would instead use the size in bytes based on actual data type.

Test Plan:
Added unit tests BatchMatMulMemCostTest:

buck test //caffe2/caffe2/fb/fbgemm:batch_matmul_op_test -- BatchMatMulMemCostTest

Extended existing unit test test_columnwise_concat for different data types:

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

Differential Revision: D30561459

fbshipit-source-id: 976fa5167097a35af548498480001aafd7851d93
caffe2/core/operator_schema.h
caffe2/operators/batch_matmul_op.cc
caffe2/operators/concat_split_op.cc
caffe2/operators/conv_pool_op_base.h
caffe2/operators/distance_op.cc
caffe2/operators/fc_inference.cc
caffe2/operators/one_hot_ops.cc
caffe2/operators/utility_ops.cc
caffe2/python/operator_test/concat_op_cost_test.py
caffe2/python/workspace_test.py
caffe2/sgd/adagrad_op.cc