mirror of https://github.com/llvm/torch-mlir
66 lines
3.3 KiB
TableGen
66 lines
3.3 KiB
TableGen
|
//===-- Passes.td - Pass definition file -------------------*- tablegen -*-===//
|
||
|
|
//
|
||
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
||
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
||
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||
|
|
//
|
||
|
|
//===----------------------------------------------------------------------===//
|
||
|
|
|
||
|
|
#ifndef NPCOMP_TORCH_PASSES
|
||
|
|
#define NPCOMP_TORCH_PASSES
|
||
|
|
|
||
|
|
include "mlir/Pass/PassBase.td"
|
||
|
|
|
||
|
|
def GlobalizeObjectGraph : Pass<"torch-globalize-object-graph", "ModuleOp"> {
|
||
|
|
let summary = "Converts TorchScript object graphs to a globalized form";
|
||
|
|
let constructor = "mlir::NPCOMP::Torch::createGlobalizeObjectGraphPass()";
|
||
|
|
let description = [{
|
||
|
|
This pass converts a subset of possible TorchScript modules into a
|
||
|
|
more restrictive lower-level form that strips away the need to be
|
||
|
|
concerned with instances of !torch.nn.Module<...> type. Specifically,
|
||
|
|
the object graph is flattened into a set of discrete globals
|
||
|
|
(`torch.global_slot`) that hold the program state.
|
||
|
|
|
||
|
|
The overarching goal is for a strict correspondence between the original
|
||
|
|
`torch.nn.Module` (call it `root`) that the user `torch.jit.script`'ed, and
|
||
|
|
the public interface of the resulting MLIR module. Specifically:
|
||
|
|
- The call `root.encoder.forward(...)` in Python corresponds to invoking
|
||
|
|
the `func @encoder.forward` on the resulting MLIR module.
|
||
|
|
- The data member access `root.decoder.ids_to_strings_table` in Python
|
||
|
|
corresponds to accessing the
|
||
|
|
`torch.global_slot @decoder.ids_to_strings_table` on the resulting
|
||
|
|
MLIR module.
|
||
|
|
In effect, the entire MLIR module corresponds to an instance of the `root`
|
||
|
|
object. This matches with the intuitive behavior desired for deployment:
|
||
|
|
When the MLIR module (or, more likely, a compiled artifact derived from it)
|
||
|
|
is loaded in a deployed environment, it is equivalent to recreating the
|
||
|
|
original `root` object.
|
||
|
|
|
||
|
|
This pass performs a complete change of the externally visible calling
|
||
|
|
convention of the MLIR module for a graph of objects and methods to a
|
||
|
|
fixed set of globals and functions.
|
||
|
|
|
||
|
|
Of course, only a subset of programs can be transformed, and this pass fails
|
||
|
|
with an error if the conditions are violated.
|
||
|
|
|
||
|
|
Specifically, the restrictions are:
|
||
|
|
- There must be a unique torch.nn_module that is not the value of a slot
|
||
|
|
of any other torch.nn_module
|
||
|
|
- Rationale: Allows us to have a notion of a unique "root" op, which is
|
||
|
|
used to define linkage. This also matches how TorchScript imports in
|
||
|
|
practice (`torch.jit.script` imports a single root object).
|
||
|
|
- There must be exactly one instance of each torch.class_type. Equivalently,
|
||
|
|
Every torch.nn_module must have a distinct type.
|
||
|
|
- Rationale: This guarantee precludes things like selecting between
|
||
|
|
multiple modules dynamically at runtime, which would require indirecting
|
||
|
|
between the separate storage of each instance.
|
||
|
|
- All torch.nn_module's must be reachable by a unique path from the root
|
||
|
|
- Rationale: Eliminates possibility of potentially exponential number of
|
||
|
|
paths. Or worse, infinite number of paths when considering cyclic
|
||
|
|
object graphs. Also as of Feb 2021, TorchScript won't import into
|
||
|
|
this form (it has a bug related to the identity of submodules).
|
||
|
|
}];
|
||
|
|
}
|
||
|
|
|
||
|
|
#endif // NPCOMP_TORCH_PASSES
|