[torchgen] Support multiple namespace in NativeFunctions.h (#79733)

Summary:
This is a follow up to #78015. This PR
* introduces namespace logic for generating `NativeFunctions.h`.
* adds helper function to extract namespace from string
* relaxes the constraint on the levels we support for custom kernel namespace to 2

Test Plan:
Yaml entry:
```
- func: unsqueeze.out(Tensor(a) self, int dim, *, Tensor(a!) out) -> Tensor(a!)
  variants: function
  device_check: NoCheck
  dispatch:
    CPU: custom_1::custom_2::unsqueeze
```

Generated `NativeFunctions.h`:

```
namespace custom_1 {
namespace custom_2 {
namespace native {
    TORCH_API at::Tensor & unsqueeze(const at::Tensor & self, int64_t dim, at::Tensor & out);
} // namespace native
} // namespace custom_2
} // namespace custom_1

```

Differential Revision: D37198111

Pull Request resolved: https://github.com/pytorch/pytorch/pull/79733
Approved by: https://github.com/bdhirsh
This commit is contained in:
Mengwei Liu
2022-07-08 21:56:49 +00:00
committed by PyTorch MergeBot
parent ff6655defb
commit 5c8a9803c8
8 changed files with 165 additions and 95 deletions

View File

@ -8,7 +8,6 @@ from typing import List, Dict, Union, Sequence, Optional
from torchgen.gen import (
get_grouped_native_functions,
parse_native_yaml,
NamespaceHelper,
)
from torchgen.model import (
BackendIndex,
@ -19,7 +18,14 @@ from torchgen.model import (
OperatorName,
)
from torchgen.selective_build.selector import SelectiveBuilder
from torchgen.utils import Target, concatMap, context, YamlLoader, FileManager
from torchgen.utils import (
Target,
concatMap,
context,
YamlLoader,
FileManager,
NamespaceHelper,
)
from torchgen.context import native_function_manager
from torchgen.code_template import CodeTemplate
import torchgen.dest as dest