Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/59466
Change saved parameter type from at::Tensor to at::IValue to support custom
class parameters, e.g. `__torch__.torch.classes.xnnpack.Conv2dOpContext`.
The NNC produced kernels won't deal with custom class parameters directly.
They simply pass through to the external operators that take these custom
class parameters, e.g. `prepacked::conv2d_clamp_run`.
It will reuse the `__getstate__` and `__setstate__` methods on the custom class
to persist and restore the state of the parameters.
When calling into the kernel, it will pass in the untyped raw pointer of the custom
class objects to the kernel as `void*`. It's similar to the regular tensor parameters,
for which it will pass in the raw data pointer of the tensor storage. The generated
kernel needs to hardcode the expected type for each parameter and cast before
calling the external ops.
ghstack-source-id: 131897904
Test Plan: - unit tests
Reviewed By: kimishpatel
Differential Revision: D28902496
fbshipit-source-id: 4b2c0895dd28f0b7d344aa08183d42ad6a355dae
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/56851
This is part of the changes to enable NNC AOT compilation for mobile.
At the end of the ahead-of-time compilation the compiler produces two sets of artifacts:
1. "compiled assembly code" - kernel functions in assembly format optimized for target platforms;
2. "compiled model" - regular TorchScript model that contains serialized parameters (weights/bias/etc) and invokes kernel functions via "handles" (name/version id/input & output specs/etc of the kernel functions).
This PR introduces a set of classes to represent kernel functions (a.k.a "handles"), which can be serialized/deserialized into/from the "compiled model" as an IValue.
Also introduces APIs to register/look-up "compiled assembly code".
ghstack-source-id: 128285802
Test Plan:
- unit tests
- for FB build environment:
buck test //caffe2/test/mobile/nnc:mobile_nnc
Reviewed By: kimishpatel, raziel
Differential Revision: D27921866
fbshipit-source-id: 4c2a4d8a4d072fc259416ae674b3b494f0ca56f3
