mirror of https://github.com/llvm/torch-mlir
c837dbb077
This primarily unlocks proper handling of free functions (that is,
functions that are not methods of any torch.nn.Module).
Recommended review order:
- `ivalue_importer.cpp` + `ivalue_import/functions*.py`
- `GlobalizeObjectGraph.cpp` + test case
- misc other stuff
The `torch::jit::CompilationUnit` is basically a backing store or
"context" holding all the possible functions in the program. The
previous code was not explicitly accessing this data structure, since it
just imported the `torch::jit::Function`'s that it saw attached to
methods.
Subtly, any time a TorchScript module called into a free function, the
free function gets incorporated into the torch::jit::CompilationUnit,
but doesn't show up anywhere when dumping the module, except in the
curious pattern:
```
%5 : Function = prim::Constant[name="adaptive_avg_pool2d"]()
%6 : Tensor = prim::CallFunction(%5, %input.1, %4)
```
That is, calls are indirect calls, and are accessed via `prim::Constant`
materializing a function object. Even stranger, the `name` attribute here
doesn't really even tell the full story -- it doesn't correspond to
anything. It turns out that the c10::FunctionType itself actually holds
a pointer to the `torch::jit::Function` in the compilation unit
directly (so there is actually no indirection in prim::CallMethod,
because any two values of the same FunctionType call the same
function!). E.g. when converting the IR to bytecode, the "name" is
ignored [code link](
|
||
---|---|---|
.. | ||
README.md | ||
add3.py | ||
debug-info.py | ||
elif.py | ||
errors.py | ||
if.py | ||
list.py | ||
prim.py | ||
tuple.py | ||
types-bool.py | ||
types-none.py |
README.md
node_import
Most of the tests in this directory test the importing of TorchScript
torch::jit::Graph
's.
However, TorchScript graphs don't really correspond directly to anything on
the MLIR side. They are a weird combination of a context, builder, and
function and just holds a torch::jit::Block
. It is torch::jit::Node
and torch::jit::Block
which form the recursive structure analogous to
MLIR's operation/region/block.
torch::jit::Node
==mlir::Operation
,torch::jit::Block
==mlir::Region
containing singlemlir::Block
Hence the name of this directory and the corresponding code in node_importer.h/cpp.