mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
efca51e171
Follows #131996 Pull Request resolved: https://github.com/pytorch/pytorch/pull/131997 Approved by: https://github.com/Skylion007
38 lines
1.3 KiB
C++
38 lines
1.3 KiB
C++
#pragma once
|
|
#include <torch/csrc/Export.h>
|
|
#include <string>
|
|
#include <unordered_map>
|
|
|
|
// `TorchScript` offers a simple optimization limit checker
|
|
// that can be configured through environment variable `PYTORCH_JIT_OPT_LIMIT`.
|
|
// The purpose is to limit how many optimization you can make per pass.
|
|
// This is useful for debugging any passes.
|
|
|
|
// Opt limit checker is enabled on a per file basis (hence per pass). For
|
|
// example, in `constant_propagation.cpp`, `PYTORCH_JIT_OPT_LIMIT` should be set
|
|
// to `constant_propagation=<opt_limit>` or, simply, to
|
|
// `constant_propagation=<opt_limit>` where <opt_limit> is the number of
|
|
// optimizations you want to make for the pass. (i.e.
|
|
// `PYTORCH_JIT_OPT_LIMIT="constant_propagation=<opt_limit>"`).
|
|
|
|
// Multiple files can be configured by separating each file name with a colon
|
|
// `:` as in the following example,
|
|
// `PYTORCH_JIT_OPT_LIMIT="constant_propagation=<opt_limit>:dead_code_elimination=<opt_limit>"`
|
|
|
|
// You can call opt limiter by calling JIT_OPT_ALLOWED. It will return true if
|
|
// we haven't reached the optimization limit yet. Otherwise, it will return
|
|
// false. Typical usage:
|
|
|
|
// if (!JIT_OPT_ALLOWED) {
|
|
// GRAPH_DUMP(...); //supplied from jit_log
|
|
// return;
|
|
// }
|
|
|
|
namespace torch::jit {
|
|
|
|
TORCH_API bool opt_limit(const char* pass_name);
|
|
|
|
#define JIT_OPT_ALLOWED opt_limit(__FILE__)
|
|
|
|
} // namespace torch::jit
|