Bin Bao
|
6066e5929f
|
[LT] Codegen ReuseNode for supported ops
Summary:
1. Update the codegen script to add a TrieCache lookup (ReuseNode)
before creating a new IR node. The following is an example generated
code,
```
at::Tensor LazyNativeFunctions::add(const at::Tensor & self, const at::Tensor & other, const at::Scalar & alpha) {
...
torch::lazy::NodePtr node = torch::lazy::ReuseNode<AddTensor>(lazy_self->GetIrValue(), lazy_other->GetIrValue(), node_alpha);
if (!node) {
auto out_meta = at::meta::add(self, other, alpha);
std::vector<Shape> shapes{Shape(out_meta.scalar_type(), out_meta.sizes().vec())};
TORCH_INTERNAL_ASSERT(shapes.size() == 1);
if(symbolicShapeEnabled()){
std::vector<jit::IValue> inputs = { self, other, alpha };
char* schema_str = "aten::add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor";
applySymbolicShapesOnLT(schema_str, inputs, shapes);
}
node = torch::lazy::MakeNode<AddTensor>(lazy_self->GetIrValue(), lazy_other->GetIrValue(), node_alpha, std::move(shapes));
CacheNode(node);
}
...
}
```
2. TrieCache lookup depends on each IR node subclass to provide its own
comparison function. The following is an example generated code,
```
bool CanBeReused(const torch::lazy::Value& self, const torch::lazy::Value& other, const torch::lazy::Value& alpha) const {
size_t i = 0;
return (operand(i++) == self &&
operand(i++) == other &&
operand(i++) == alpha);
}
```
3. DeviceData is specially handled.
4. Non-codegen op changes are coming a separate PR.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/76738
Approved by: https://github.com/JackCaoG, https://github.com/wconstab
|
2022-05-13 19:13:58 +00:00 |
|