mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
6066e5929f
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
|
||
|---|---|---|
| .. | ||
| init.cpp | ||
| init.h | ||
| python_util.cpp | ||
| python_util.h | ||
| README.md | ||
Lazy Tensor Python Code
Lazy Tensor Core is part of libtorch, which can not depend on python.
Parts of lazy tensor core use python for 2 purposes A) py bindings let python programs call into lazy tensor c++ code B) lazy tensor core calls into python to use it (e.g. for grabbing stack traces)
(A) is trivial since the python bindings only depend on libtorch; (B) requires making libtorch_python register a function with libtorch if loaded, and having a default (no-op) function otherwise. Any functionality that strictly needs to depend on python should be part of the 'python' folder.