mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
3842140fd5
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/72730
This diff contains changes from several PRs landed to lazy_tensor_staging branch.
- generating 'fallback' overrides for each codegenned op, useful for debugging
- supports operators which are missing aten:: symbols for op names, instead using their string counterpart
- makes the IR class a base class instead of hardcoding the assumption of TS
Test Plan: tested on lazy_tensor_staging branch
Reviewed By: desertfire
Differential Revision: D34178476
fbshipit-source-id: 7190b2e0d82b4eb1f4510c858c24446c6df3f9d0
(cherry picked from commit 6713d3f0ef)
84 lines
2.1 KiB
C++
84 lines
2.1 KiB
C++
#include <gtest/gtest.h>
|
|
|
|
#include <c10/util/Exception.h>
|
|
#include <torch/csrc/lazy/core/config.h>
|
|
#include <torch/csrc/lazy/core/ir.h>
|
|
#include <torch/csrc/lazy/core/ir_metadata.h>
|
|
#include <torch/csrc/lazy/core/ir_util.h>
|
|
|
|
namespace torch {
|
|
namespace lazy {
|
|
|
|
class IrUtilNode : public Node {
|
|
public:
|
|
explicit IrUtilNode()
|
|
: Node(OpKind(), /* num_outputs */ 1, /* hash_func */ [&](bool bakeInSizes) -> hash_t { return Hash(0); }) {}
|
|
~IrUtilNode() override = default;
|
|
|
|
void AddOperand(Value v) {
|
|
if (!v.node) {
|
|
return;
|
|
}
|
|
operands_as_outputs_.emplace_back(v.node.get(), v.index);
|
|
operands_.push_back(std::move(v.node));
|
|
}
|
|
|
|
const std::vector<Output>& operands() const override {
|
|
return operands_as_outputs_;
|
|
}
|
|
|
|
const Output& operand(size_t i) const override {
|
|
return operands_as_outputs_.at(i);
|
|
}
|
|
|
|
private:
|
|
std::vector<NodePtr> operands_;
|
|
std::vector<Output> operands_as_outputs_;
|
|
};
|
|
|
|
/* a
|
|
* / \
|
|
*b c
|
|
* \ /
|
|
* d
|
|
* Post-order: d c b a
|
|
*/
|
|
TEST(IrUtilTest, BasicTest) {
|
|
NodePtr a = MakeNode<IrUtilNode>();
|
|
NodePtr b = MakeNode<IrUtilNode>();
|
|
NodePtr c = MakeNode<IrUtilNode>();
|
|
NodePtr d = MakeNode<IrUtilNode>();
|
|
|
|
dynamic_cast<IrUtilNode*>(a.get())->AddOperand(Value(b, 0));
|
|
dynamic_cast<IrUtilNode*>(a.get())->AddOperand(Value(c, 1));
|
|
dynamic_cast<IrUtilNode*>(b.get())->AddOperand(Value(d, 0));
|
|
dynamic_cast<IrUtilNode*>(c.get())->AddOperand(Value(d, 0));
|
|
|
|
std::vector<Node*> postorder = Util::ComputePostOrder({a.get()});
|
|
EXPECT_EQ(postorder.size(), 4);
|
|
EXPECT_EQ(postorder.at(0), d.get());
|
|
EXPECT_EQ(postorder.at(1), c.get());
|
|
EXPECT_EQ(postorder.at(2), b.get());
|
|
EXPECT_EQ(postorder.at(3), a.get());
|
|
}
|
|
|
|
/* a
|
|
* / \
|
|
*b---c
|
|
* Post-order: not valid
|
|
*/
|
|
TEST(IrUtilTest, TestCircle) {
|
|
NodePtr a = MakeNode<IrUtilNode>();
|
|
NodePtr b = MakeNode<IrUtilNode>();
|
|
NodePtr c = MakeNode<IrUtilNode>();
|
|
|
|
dynamic_cast<IrUtilNode*>(a.get())->AddOperand(Value(b, 0));
|
|
dynamic_cast<IrUtilNode*>(b.get())->AddOperand(Value(c, 0));
|
|
dynamic_cast<IrUtilNode*>(c.get())->AddOperand(Value(a, 0));
|
|
|
|
EXPECT_THROW(Util::ComputePostOrder({a.get()}), c10::Error);
|
|
}
|
|
|
|
} // namespace lazy
|
|
} // namespace torch
|