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
107
108
109
|
#include <torch/csrc/jit/codegen/fuser/interface.h>
#include <torch/csrc/jit/codegen/fuser/compiler.h>
#include <torch/csrc/jit/codegen/fuser/executor.h>
#include <torch/csrc/jit/codegen/fuser/fallback.h>
#include <torch/csrc/jit/codegen/fuser/kernel_cache.h>
#include <c10/util/Flags.h>
#include <stdexcept>
namespace torch {
namespace jit {
namespace detail {
#ifdef TORCH_ENABLE_LLVM
bool cpu_fuser_enabled = true;
#else
bool cpu_fuser_enabled = false;
#endif
// note: this doesn't necessarily enable NNC because NVFuser might override it
bool gpu_fuser_enabled = true;
} // namespace detail
int64_t registerFusion(const Node* fusion_group) {
return fuser::registerFusion(fusion_group);
}
void runFusion(const int64_t key, Stack& stack) {
const auto result = fuser::runFusion(key, stack);
if (!result)
fuser::runFallback(key, stack);
}
bool canFuseOnCPU() {
return fuser::hasFusionBackend(DeviceType::CPU) && detail::cpu_fuser_enabled;
}
bool canFuseOnGPU() {
return fuser::hasFusionBackend(DeviceType::CUDA) && detail::gpu_fuser_enabled;
}
void overrideCanFuseOnCPU(bool value) {
detail::cpu_fuser_enabled = value;
}
void overrideCanFuseOnGPU(bool value) {
detail::gpu_fuser_enabled = value;
}
// Uses the above interface by stuffing the graph into a node and treating that
// node as a fusion group.
std::vector<at::Tensor> debugLaunchGraph(
Graph& graph,
at::ArrayRef<at::Tensor> inputs) {
// Creates a fusion group node
auto wrapper_graph = std::make_shared<Graph>();
Node* fusion_group = wrapper_graph->insertNode(
wrapper_graph->createWithSubgraph(prim::FusionGroup));
fusion_group->g_(attr::Subgraph, graph.copy());
for (size_t i = 0; i < graph.inputs().size(); ++i) {
fusion_group->addInput(wrapper_graph->addInput());
}
for (size_t i = 0; i < graph.outputs().size(); ++i) {
wrapper_graph->registerOutput(fusion_group->addOutput());
}
// Creates the stack, registers and runs the fusion
Stack stack = fmap<IValue>(inputs);
const auto key = fuser::registerFusion(fusion_group);
fuser::runFusion(key, stack);
return fmap(stack, [](const IValue& iv) { return iv.toTensor(); });
}
std::string debugGetFusedKernelCode(
Graph& graph,
at::ArrayRef<at::Tensor> inputs) {
// Creates a fusion group node
auto wrapper_graph = std::make_shared<Graph>();
Node* fusion_group = wrapper_graph->insertNode(
wrapper_graph->createWithSubgraph(prim::FusionGroup));
fusion_group->g_(attr::Subgraph, graph.copy());
for (size_t i = 0; i < graph.inputs().size(); ++i) {
fusion_group->addInput(wrapper_graph->addInput());
}
for (size_t i = 0; i < graph.outputs().size(); ++i) {
wrapper_graph->registerOutput(fusion_group->addOutput());
}
// Creates the stack, registers and runs the fusion
Stack stack = fmap<IValue>(inputs);
const auto key = fuser::registerFusion(fusion_group);
std::string code;
if (!fuser::runFusion(key, stack, &code)) {
throw std::runtime_error("Could not run fusion for graph");
}
return code;
}
size_t nCompiledKernels() {
return fuser::nCompiledKernels();
}
} // namespace jit
} // namespace torch
|
