intel-extension-for-pytorch
73 строки · 2.4 Кб
1#include "initialization.h"
2#include <oneapi/dnnl/dnnl_graph.hpp>
3#include <stdio.h>
4#include <torch/csrc/jit/passes/autocast.h>
5#include <torch/csrc/jit/passes/pass_manager.h>
6#include <torch/torch.h>
7#include <regex>
8#include <string>
9#include "auto_opt_config.h"
10#include "fusion_pass.h"
11#include "version.h"
12
13namespace torch_ipex {
14
15void init_jit_fusion_pass() {
16// jit fusion pass
17torch::jit::registerPrePass([](std::shared_ptr<torch::jit::Graph>& g) {
18if (AutoOptConfig::singleton().get_jit_fuse()) {
19torch_ipex::jit::FusionPass(g);
20}
21});
22}
23
24void disable_autocast_for_jit_script() {
25// We need disable autocast pass by default after
26// https://github.com/pytorch/pytorch/pull/74178. Will remove this after we
27// can extend the cast policy for this autocast pass.
28torch::jit::setAutocastMode(false);
29}
30
31void enable_onednn_graph_constant_tensor_cache() {
32// oneDNN graph constant cache was enabled by default.
33// It is changed to be disabled by default now to avoid potential OOM issues.
34// Turn it on here to enable weight cache in inference to avoid performance
35// regressions.
36dnnl::graph::set_constant_tensor_cache(true);
37}
38
39InitIPEX::InitIPEX() = default;
40InitIPEX::~InitIPEX() = default;
41InitIPEX::InitIPEX(InitIPEX&&) noexcept = default;
42
43static auto init = InitIPEX()
44.init(&init_jit_fusion_pass)
45.init(&disable_autocast_for_jit_script)
46.init(&enable_onednn_graph_constant_tensor_cache);
47
48void InitIPEX::check_pytorch_version() {
49int IPEX_VERSION_MAJOR = 0;
50int IPEX_VERSION_MINOR = 0;
51const std::regex regex("(\\d+)\\.(\\d+).*");
52const std::string ipex_version = torch_ipex::__version__();
53std::smatch match;
54if (std::regex_match(ipex_version, match, regex)) {
55if (match.size() == 3) {
56IPEX_VERSION_MAJOR = std::stoi(match[1]);
57IPEX_VERSION_MINOR = std::stoi(match[2]);
58}
59}
60if (IPEX_VERSION_MAJOR != TORCH_VERSION_MAJOR ||
61IPEX_VERSION_MINOR != TORCH_VERSION_MINOR) {
62printf(
63"ERROR! Intel® Extension for PyTorch* needs to work with PyTorch/libtorch %d.%d.*, but PyTorch/libtorch %d.%d.%d is found. Please switch to the matching version and run again.\n",
64IPEX_VERSION_MAJOR,
65IPEX_VERSION_MINOR,
66TORCH_VERSION_MAJOR,
67TORCH_VERSION_MINOR,
68TORCH_VERSION_PATCH);
69exit(127);
70}
71}
72
73} // namespace torch_ipex
74