_(FromQualString) \
_(InternedStrings) \
_(IValue) \
+ _(PassManagement) \
_(Proto) \
_(RegisterFusionCachesKernel) \
_(SchemaParser) \
#include "torch/csrc/jit/fuser/interface.h"
#include "torch/csrc/jit/import.h"
#include "torch/csrc/jit/interpreter.h"
+#include "torch/csrc/jit/pass_manager.h"
#include "torch/csrc/jit/passes/alias_analysis.h"
#include "torch/csrc/jit/passes/common_subexpression_elimination.h"
#include "torch/csrc/jit/passes/constant_propagation.h"
AT_ASSERT(result.toTensor().item<float>() == 6)
}
+static int testPassValue = 0;
+void fakePass(std::shared_ptr<Graph>& g) {
+ testPassValue++;
+ return;
+}
+
+RegisterPass p(fakePass);
+
+void testPassManagement() {
+ std::shared_ptr<Graph> graph = std::make_shared<Graph>();
+ script::parseIR(
+ R"IR(
+graph(%a):
+ return (%a))IR",
+ &*graph);
+
+ std::vector<IValue> stack = {IValue(torch::randn({22}, at::kCPU))};
+ auto run = [&](std::shared_ptr<Graph>& graph, std::vector<IValue> stack) {
+ GraphExecutor executor(graph);
+ executor.run(stack);
+ return stack;
+ };
+ run(graph, stack);
+ AT_ASSERT(testPassValue);
+}
+
} // namespace test
} // namespace jit
} // namespace torch
"torch/csrc/jit/constants.cpp",
"torch/csrc/jit/node_hashing.cpp",
"torch/csrc/jit/export.cpp",
+ "torch/csrc/jit/pass_manager.cpp",
"torch/csrc/jit/pickler.cpp",
"torch/csrc/jit/graph_executor.cpp",
"torch/csrc/jit/import.cpp",
${TORCH_SRC_DIR}/csrc/jit/attributes.cpp
${TORCH_SRC_DIR}/csrc/jit/argument_spec.cpp
${TORCH_SRC_DIR}/csrc/jit/export.cpp
+ ${TORCH_SRC_DIR}/csrc/jit/pass_manager.cpp
${TORCH_SRC_DIR}/csrc/jit/pickler.cpp
${TORCH_SRC_DIR}/csrc/jit/generated/register_aten_ops_0.cpp
${TORCH_SRC_DIR}/csrc/jit/generated/register_aten_ops_1.cpp
#include <torch/csrc/jit/custom_operator.h>
#include <torch/csrc/jit/interpreter.h>
#include <torch/csrc/jit/ir.h>
+#include <torch/csrc/jit/pass_manager.h>
#include <torch/csrc/jit/passes/batch_mm.h>
#include <torch/csrc/jit/passes/canonicalize_ops.h>
#include <torch/csrc/jit/passes/common_subexpression_elimination.h>
}
void runNondiffOptimization(std::shared_ptr<Graph>& graph) {
+ for (const auto& pass : getCustomPasses()) {
+ pass(graph);
+ }
FuseGraph(graph);
}
--- /dev/null
+#include <torch/csrc/jit/pass_manager.h>
+
+namespace torch {
+namespace jit {
+
+std::vector<Pass>& getCustomPasses() {
+ static std::vector<Pass> passes;
+ return passes;
+}
+
+RegisterPass::RegisterPass(Pass p) {
+ getCustomPasses().emplace_back(std::move(p));
+}
+
+} // namespace jit
+} // namespace torch
--- /dev/null
+#pragma once
+
+#include <torch/csrc/jit/ir.h>
+
+/* `getCustomPasses()` returns a vector of passes that will be executed after
+ * differentiation but before any fusion. This is the de-facto location
+ * for compiler backends to insert passes.
+ *
+ * Static registration of a pass can be done by creating a global
+ * `RegisterPass r(Pass)` variable in a compilation unit.
+ *
+ * pass_manager.h uses a Meyer's singleton
+ * to store a vector of `Pass`es, which modify the IR graph in place.
+ */
+
+namespace torch {
+namespace jit {
+
+// A pass modifies a Graph in place.
+using Pass = std::function<void(std::shared_ptr<Graph>&)>;
+
+std::vector<Pass>& getCustomPasses();
+
+struct RegisterPass {
+ RegisterPass(Pass p);
+};
+
+} // namespace jit
+} // namespace torch