1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
|
#pragma once
#include <torch/csrc/lazy/ts_backend/ts_node.h>
#include <torch/csrc/lazy/core/ir_builder.h>
namespace torch {
namespace lazy {
// Generic IR Node implementation for nodes which can simply be described by a
// specific OpKind and a lowering function. IR nodes carrying
// metadata should not be using this class TORCH_API (and have the metadata
// captured by the LowerFn), but they should instead create a dedicated IR node.
// Doing the former would limit IR introspection.
class TORCH_API Generic : public TsNode {
public:
Generic(
OpKind op,
OpList operands,
Shape shape,
size_t num_outputs = 1,
hash_t hash_seed = static_cast<uint32_t>(0x5a2d296e9));
Generic(
OpKind op,
OpList operands,
const std::function<Shape()>& shape_fn,
size_t num_outputs = 1,
hash_t hash_seed = static_cast<uint32_t>(0x5a2d296e9));
Generic(
OpKind op,
OpList operands,
size_t num_outputs = 1,
hash_t hash_seed = static_cast<uint32_t>(0x5a2d296e9));
Generic(OpKind op, Shape shape, size_t num_outputs, hash_t hash_seed);
private:
hash_t hash_seed_;
};
inline NodePtr GenericOp(
OpKind op,
OpList operands,
Shape shape,
size_t num_outputs = 1,
hash_t hash_seed = static_cast<uint32_t>(0x5a2d296e9)) {
return MakeNode<Generic>(
op, operands, std::move(shape), num_outputs, hash_seed);
}
} // namespace lazy
} // namespace torch
|
