[structured] Preserve computed elements from meta func to impl (#61746)
authorMeghan Lele <meghanl@fb.com>
Wed, 1 Sep 2021 21:24:54 +0000 (14:24 -0700)
committerFacebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
Wed, 1 Sep 2021 21:34:25 +0000 (14:34 -0700)
commit968d7ee46a66316557a04a64333d2810b544ed9b
treefff6ee1432aac326854ab5acbd1389ddfe961795
parent4aad366111f88010c9e4027d054da2fe357e1bfe
[structured] Preserve computed elements from meta func to impl (#61746)

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

**Summary**
This commit introduces a new feature for structured kernels that allows
kernels to declare quantities as "precomputed" in
`native_functions.yaml`, compute them once in the `meta` function and
reuse them again in the `impl`. The names and types of these quantities
are used to generate code for a struct containing them that the `meta`
function must return. In the case of a handful of surveyed kernels
(`all,`, `any`, `avg_pool2d`), these quantities that are used both in
the `meta` and `impl` have the same meaning as certain kernel arguments
and in fact supersede them. Accordingly, the correspondence between a
kernel argument and the precomputed elements that supersede it is also
captured in `native_functions.yaml`. This information is used to unpack
the struct returned by `meta` and pass its contents correctly to the
`impl` function.

The primary goal is to avoid recompute and enhance developer experience
(e.g. sometimes people can forget to compute these elements while
porting a kernel).

Test Plan: Imported from OSS

Reviewed By: tugsbayasgalan

Differential Revision: D30407831

Pulled By: SplitInfinity

fbshipit-source-id: 00975525ea373721fe52d06f75cd4ac91f3dc556
aten/src/ATen/TensorMeta.h
aten/src/ATen/native/AveragePool2d.cpp
aten/src/ATen/native/Pool.h
aten/src/ATen/native/ReduceOps.cpp
aten/src/ATen/native/cpu/AvgPoolKernel.cpp
aten/src/ATen/native/cuda/AveragePool2d.cu
aten/src/ATen/native/native_functions.yaml
tools/codegen/api/structured.py
tools/codegen/dest/register_dispatch_key.py
tools/codegen/gen.py
tools/codegen/model.py