pytorch

5e6f2966 - Structured Kernel Precompute codegen handle fields without replacement (#71368)

Structured Kernel Precompute codegen handle fields without replacement (#71368) Summary: I've added the parsing of an optional first line in native_functions.yaml after the precomputed keyword for arguments that will be precomputed without replacement. This line is optional, must be the first and does not contain any arrow. These new fields are precomputed as before in the meta function and added to the precompute struct returned by the meta function. For now I've put them as last args of the impl function where they can be reused. example: native_function.yaml: ``` ... precomputed: - int numBatch, int numPlanes, int inputT, int inputH, int inputW <- new - kernel_size -> int poolSizeT, int poolSizeH, int poolSizeW - output_size -> int outputT, int outputH, int outputW ``` meta: ``` TORCH_PRECOMPUTE_META_FUNC(fractional_max_pool3d)( const at::Tensor& input_, IntArrayRef pool_size, IntArrayRef output_size, const at::Tensor& randomSamples ) { ... return TORCH_PRECOMPUTE_STRUCT(fractional_max_pool3d)().set_numBatch(numBatch).set_numPlanes(numPlanes).set_inputT(inputT).set_inputH(inputH).set_inputW(inputW) .set_poolSizeT(poolSizeT) ... } ``` impl: ``` TORCH_IMPL_FUNC(fractional_max_pool3d_out_cpu)( const at::Tensor& input_, int64_t poolSizeT, int64_t poolSizeH, int64_t poolSizeW, int64_t outputT, int64_t outputH, int64_t outputW, const at::Tensor& randomSamples, const at::Tensor& output, const at::Tensor& indices, int64_t numBatch, <- for now I've put them here int64_t numPlanes, int64_t inputT, int64_t inputH, int64_t inputW) { ``` Fixes https://github.com/pytorch/pytorch/issues/71314 Pull Request resolved: https://github.com/pytorch/pytorch/pull/71368 Reviewed By: zou3519 Differential Revision: D33683984 Pulled By: bdhirsh fbshipit-source-id: 33066dd92b8743aadf0dc8102f6bf0689f843242 (cherry picked from commit 64e46af6a4e7dcd28bd02ae3ab3c6d979b4caa40)