-
Notifications
You must be signed in to change notification settings - Fork 81
Expand file tree
/
Copy pathinternal_ir.cpp
More file actions
106 lines (91 loc) · 2.57 KB
/
internal_ir.cpp
File metadata and controls
106 lines (91 loc) · 2.57 KB
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
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
// clang-format off
/*
* SPDX-FileCopyrightText: Copyright (c) 2025-present NVIDIA CORPORATION & AFFILIATES.
* All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*/
// clang-format on
#include <bindings.h>
#include <ir/internal_nodes.h>
namespace nvfuser::python {
// For all nodes, use multiple inheritance to disable destructor with
// std::unique_ptr<Statement, py::nodelete>. This class will
// disable memory management because it is handled automatically by IrContainer.
namespace {
void bindInternalFusionIr(py::module& nvfuser) {
py::class_<Split, Expr, std::unique_ptr<Split, py::nodelete>> split(
nvfuser, "Split");
split.def(
"__str__",
[](Split* self) { return self->toString(/*indent_size=*/0); },
"Convert the Split to a string representation.");
split.def(
"outer",
&Split::outer,
R"(
Get the outer of this Split.
Returns
-------
IterDomain
The outer of this Split.
)");
split.def(
"inner",
&Split::inner,
R"(
Get the inner of this Split.
Returns
-------
IterDomain
The inner of this Split.
)");
py::class_<Merge, Expr, std::unique_ptr<Merge, py::nodelete>> merge(
nvfuser, "Merge");
merge.def(
"__str__",
[](Merge* self) { return self->toString(/*indent_size=*/0); },
"Convert the Merge to a string representation.");
merge.def(
"outer",
&Merge::outer,
R"(
Get the outer of this Merge.
Returns
-------
IterDomain
The outer of this Merge.
)");
merge.def(
"inner",
&Merge::inner,
R"(
Get the inner of this Merge.
Returns
-------
IterDomain
The inner of this Merge.
)");
py::class_<BroadcastOp, Expr, std::unique_ptr<BroadcastOp, py::nodelete>>
broadcast(nvfuser, "BroadcastOp");
broadcast.def(
"__str__",
[](BroadcastOp* self) { return self->toString(/*indent_size=*/0); },
"Convert the BroadcastOp to a string representation.");
py::class_<ReshapeOp, Expr, std::unique_ptr<ReshapeOp, py::nodelete>> reshape(
nvfuser, "ReshapeOp");
reshape.def(
"__str__",
[](ReshapeOp* self) { return self->toString(/*indent_size=*/0); },
"Convert the ReshapeOp to a string representation.");
py::class_<SqueezeOp, Expr, std::unique_ptr<SqueezeOp, py::nodelete>> squeeze(
nvfuser, "SqueezeOp");
squeeze.def(
"__str__",
[](SqueezeOp* self) { return self->toString(/*indent_size=*/0); },
"Convert the SqueezeOp to a string representation.");
}
} // namespace
void bindInternalIr(py::module& nvfuser) {
bindInternalFusionIr(nvfuser);
}
} // namespace nvfuser::python