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
|
#pragma once
#include <oneapi/dnnl/dnnl_graph.hpp>
#include <torch/csrc/jit/codegen/onednn/operator.h>
#include <torch/csrc/jit/ir/ir.h>
namespace torch {
namespace jit {
namespace fuser {
namespace onednn {
struct OpPartitionMap {
void add(uint64_t opId, uint64_t partitionId) {
opmap_[opId] = partitionId;
}
void add(Node* n, uint64_t partitionId) {
add(Operator::getId(n), partitionId);
}
bool has(uint64_t opId) {
return opmap_.count(opId) > 0;
}
bool has(Node* n) {
return has(Operator::getId(n));
}
uint64_t get(uint64_t opId) {
return opmap_[opId];
}
uint64_t get(Node* n) {
auto opId = Operator::getId(n);
TORCH_CHECK(
has(opId),
"Node ",
n->kind().toQualString(),
" does not belong to any LLGA partition");
return get(opId);
}
private:
std::unordered_map<uint64_t, uint64_t> opmap_;
};
class LlgaGraphHelper {
public:
LlgaGraphHelper(
const std::shared_ptr<Graph>& graph,
dnnl::graph::partition::policy policy =
dnnl::graph::partition::policy::fusion);
bool shouldMerge(Node* toMerge, Node* subgraph);
bool shouldConsiderForMerge(Node* node);
bool checkForSingleOpPartition(Node* node);
Node* createSingletonSubgraph(Node* n, AliasDb& db);
void mergeNodeIntoSubgraph(Node* toMerge, Node* subgraphNode, AliasDb& db);
void unmergeIfAnyNodeIsMissing(Node* subgraphNode);
static bool isLlgaSubgraph(const Node* node);
std::vector<dnnl::graph::partition> getPartitions() const;
std::map<size_t, Value*> getTensorIdToValue() const;
private:
size_t countSupportedOps(const std::shared_ptr<Graph>& graph) const;
OpPartitionMap opToOwningPartition_;
std::vector<dnnl::graph::partition> partitions_;
std::map<size_t, Value*>
tensorIdToValue_; // map from tensorId to torch::jit::Value
};
class LlgaNodeWrapper {
public:
LlgaNodeWrapper(const Node* node);
void setOpaqueLayout(size_t offset);
bool useOpaqueLayout(size_t offset) const;
friend class LlgaGraphHelper;
private:
void initOutputLayouts();
Node* n;
};
} // namespace onednn
} // namespace fuser
} // namespace jit
} // namespace torch
|