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
|
#include <torch/csrc/jit/codegen/onednn/guard_shape.h>
#include <torch/csrc/jit/jit_log.h>
#include <torch/csrc/jit/passes/tensorexpr_fuser.h>
#include <torch/csrc/jit/passes/utils/subgraph_utils.h>
#include <torch/csrc/jit/runtime/graph_executor.h>
namespace torch {
namespace jit {
namespace fuser {
namespace onednn {
//! [ Note -- prepareFusionGroupAndGuardOutputs implementation ]
//! shamelessly copying code from NNC (tensorexpr_fuser) with very little
//! modification, original code at:
//! `torch/csrc/jit/passes/tensorexpr_fuser.cpp:prepareFusionGroupAndGuardOutputs`
//!
//! We have the assumption that LLGA does not have operators
//! depending on the content of the tensor.
void prepareFusionGroupAndGuardOutputs(Block* block) {
std::vector<Node*> fusion_groups;
for (Node* n : block->nodes()) {
for (Block* b : n->blocks()) {
prepareFusionGroupAndGuardOutputs(b);
}
if (n->kind() == prim::oneDNNFusionGroup) {
fusion_groups.push_back(n);
}
}
for (Node* fusion_group : fusion_groups) {
// TODO: add further optimization pass to removeOutputsUsedOnlyInSize,
// refer to
// `torch/csrc/jit/passes/tensorexpr_fuser.cpp:removeOutputsUsedOnlyInSize`
// removeOutputsUsedOnlyInSize(fusion_group);
insertTypeGuard(
fusion_group,
[](const TensorTypePtr& t) { return t; },
prim::oneDNNFusionGuard);
}
}
} // namespace onednn
} // namespace fuser
} // namespace jit
} // namespace torch
|
