llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity
torch-mlir/frontends/pytorch/csrc/builder/ivalue_importer.cpp

543 lines
22 KiB
C++
Raw Normal View History

Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
//===- ivalue_importer.cpp ------------------------------------------------===//
//
// This file is licensed under a pytorch-style license
// See frontends/pytorch/LICENSE for license information.
//
//===----------------------------------------------------------------------===//
#include "ivalue_importer.h"
Add ability to annotate TorchScript classes. The first use case is to annotate certain program constructs as either exported or private. In this commit we plumb it down to GlobalizeObjectGraph which makes use of this information. Recommended review order: 1. class_annotator.h/.cpp + `test/module_import/annotations/*` - New abstractions to communicate with Python code and annotate. 2. IR changes in TorchOps.td - Adding "private" attribute to various things. 3. ivalue_import.cpp changes - Module + ClassAnnotator = annotated IR 4. GlobalizeObjectGraph.cpp + tests - use new "private" attributes to create "private" IR. - also, tweak some of the op deleting mechanics, which was triggering some memory errors / assertions With this, we can run the classifier through and inline it as follows: ``` frontends/pytorch/utils/pt_util.py --import --exported-name forward ~/tmp/classifier.pt \ | npcomp-opt -torch-globalize-object-graph -inline ``` IR: https://gist.github.com/silvasean/32dcad9f6270557f412094a77cecdd69
2021-02-20 08:21:21 +08:00
#include "class_annotator.h"
Properly model "derefinement". In terms of IR structure, TorchScript allows types to vary in many circumstances where MLIR requires pointer-identical types. In particular, it is valid to pass any subtype in place of a type. For example, if an `Optional[int]` is required somewhere in the IR, it is legal to pass a value of just `int` (but not the other way around; see `torch.prim.unchecked_cast`). In effect, every *use* can have a different type. We introduce a new op `torch.derefine` that models that impedance mismatch. This op allows casting a value from one type to a type that it is a subtype of to model this behavior. Recommended review order: - TorchOps.td for new torch.derefine (and updated docs for `torch.prim.unchecked_cast`) - new test code in if.py, loop.py, function-derefine.py - new code in node_importer.cpp for handling derefinement insertion - function_importer.cpp and utils changes in torch_to_mlir_utils.cpp Properly handling derefinement on function boundaries required relayering the code so that graph_importer.cpp/.h is now function_importer.cpp/.h because only the `torch::jit::Function` (actually the `c10::FunctionSchema` it holds) knows the derefined types that are actually needed at the boundary (see `function-derefine.py` for a test). Annoyingly, this churns all the functions which are now prefixed with `__torch__.` but that is more correct anyway (that is their linkage name in the `torch::jit::CompilationUnit`; the previous `mb.import_function` was actually buggy in the case of functions calling each other as it would reference their unqualified name). With this change, we can import `resnet18` from `torchvision` :) IR: https://gist.github.com/silvasean/6426a5272d8a6c7caae533fce05ab704
2021-03-02 09:24:15 +08:00
#include "function_importer.h"
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
#include <unordered_map>
#include "mlir_utils.h"
#include "mlir-c/BuiltinAttributes.h"
#include "mlir-c/BuiltinTypes.h"
#include "mlir-c/Diagnostics.h"
#include "npcomp-c/Types.h"
Shore up error reporting for TorchScript import. This code was not exception safe -- it would leave an operation unattached to anything, which breaks MLIR's C++ data structure invariants (e.g. it cannot safely erase ops). Also, print out both the exception and any diagnostics, since they can both contain useful information.
2021-05-19 03:48:22 +08:00
#include "caffe2/core/scope_guard.h"
Get simple quantized model importing. This is enough to import the program and get it through the compilation pipeline. It of course fails at the VerifyBackendContract pass since there is a lot missing, but the final IR for a simple quantized MLP is looking pretty decent already: [IR](https://gist.github.com/silvasean/f76bccd76e9b193d396cfb2f9a11f54d) Main changes: - Add support for importing torch quantized tensors, including `torch.per_tensor_affine.create` op and `!torch.qint8` element type. - Add support for importing `LinearPackedParamsBase` (basically a weight + optional bias, but requires `torch.linear_params.create` op + `!torch.LinearParams` type to model it). This was less painful than I expected, as it has the necessary methods to opaquely unpack itself. I factored things so it should be easy to extend to other custom classes like `ConvPackedParamsBase`. - Add minimal boilerplate for importing `quantized::*` ops, with `quantized::linear` being a motivating example. - Add e2e test with simple quantized MLP (courtesy of @phoenix-meadowlark). This is somewhat of an abuse of `!numpy.ndarray` / `tensor`, as really the proper semantics of `!torch.qint8` dtype on a Torch tensor is "check the quantizer object of the tensor for side data (scale/offset, possibly per-channel) that defines the full semantics of the tensor". We don't have any such notion of "side data" for `!numpy.ndarray` / `tensor`, let alone anything that would have the associated behavior of keying off the dtype to determine if the side data is present. This will be fixed by a proper `!torch.tensor` type.
2021-05-20 02:40:48 +08:00
#include "ATen/native/quantized/cpu/packed_params.h"
Shore up error reporting for TorchScript import. This code was not exception safe -- it would leave an operation unattached to anything, which breaks MLIR's C++ data structure invariants (e.g. it cannot safely erase ops). Also, print out both the exception and any diagnostics, since they can both contain useful information.
2021-05-19 03:48:22 +08:00
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
using namespace torch_mlir;
Handle object identity correctly. This required some careful considerations when defining object identity for tensors. See the comments for how we do it. This also tracks some basic information for diagnostics.
2021-02-06 09:57:38 +08:00
// Hashing functionality for IValue's.
//
// What we want here is a strict object identity hash. This is different from
// what Python usually treats as hashing, which is a deep equality hash. In
// Python terms, what we want here is a hash of `id(x)` -- unfortunately, IValue
// is not uniformly heap allocated the way a `PyObject*` is, so special handling
// is needed. At the time of this writing, there seem to be two different
// implementations, neither of which is exactly what we want.
//
// - c10::IValue::hash static method
// - Problem: Doesn't handle certain data types, in particular objects (which
// modules are a special case of) and lists/dicts. This makes sense when
// reflecting the Python semantics.
// - c10::WeakIValue::hash method
// - Problem: it literally just returns the bits of the "union" as an int.
// This seems to read uninitialized bits for the bool variant.
//
// We use the `c10::IValue::hash` static method with special cases for data
// types that need their identity to be handled specially. `c10::IValue::hash`
// seems to be implemented in a principled way following the Python semantics,
// which is compatible with the semantics we want (for the subset it doesn't
// throw an error on).
namespace {
struct IValueHasher {
size_t operator()(const c10::IValue &ivalue) const {
if (ivalue.isObject() || ivalue.isList()) {
return std::hash<const void *>()(
static_cast<const void *>(ivalue.internalToPointer()));
}
return c10::IValue::hash(ivalue);
}
};
} // namespace
// TODO: The implementation of isSameIdentity looks vulnerable to malloc reusing
// the same memory block (if this hash function is used in an online setting,
// such as when tracing). Can we do better?
namespace {
struct IValueEq {
bool operator()(const c10::IValue &lhs, const c10::IValue &rhs) const {
return lhs.isSameIdentity(rhs);
}
};
} // namespace
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
namespace {
/// Helper class for holding state during recursive IValue import.
///
/// The intended usage pattern of this class is to construct it then call
Handle object identity correctly. This required some careful considerations when defining object identity for tensors. See the comments for how we do it. This also tracks some basic information for diagnostics.
2021-02-06 09:57:38 +08:00
/// `importIValue`.
///
/// The `importIValue` method can be called more than once, and values are
/// unified *by object identity*. For types isomorphic to Python builtin types
/// the behavior is what you would expect from `id(x)`.
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
///
Handle object identity correctly. This required some careful considerations when defining object identity for tensors. See the comments for how we do it. This also tracks some basic information for diagnostics.
2021-02-06 09:57:38 +08:00
/// For tensors, object identity is a little tricky. As background, at::Tensor
/// basically has 4 parts:
/// - at::Tensor which is a smart pointer to at::TensorImpl
/// - at::TensorImpl which holds sizes/strides/etc. and points to at::Storage
/// - the address of the at::TensorImpl is the identity of the tensor.
/// - at::Storage which is a smart pointer to at::StorageImpl
/// - at::StorageImpl which is a low-level buffer
/// - the address of the at::StorageImpl is the identity of the "storage".
///
/// Multiple different tensors can share the same underlying storage. We
/// currently import tensors by identity and emit errors in the case of tensors
/// with different identity but sharing the same storage. This is done because
/// correctly modeling the many ways that tensors can overlap and alias when
/// they share storage is difficult. Example hard cases are weird
/// strides/offsets that overlap, and even cases where the data types mismatch
/// (PyTorch allows this!).
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
class IValueImporter {
public:
Add ability to annotate TorchScript classes. The first use case is to annotate certain program constructs as either exported or private. In this commit we plumb it down to GlobalizeObjectGraph which makes use of this information. Recommended review order: 1. class_annotator.h/.cpp + `test/module_import/annotations/*` - New abstractions to communicate with Python code and annotate. 2. IR changes in TorchOps.td - Adding "private" attribute to various things. 3. ivalue_import.cpp changes - Module + ClassAnnotator = annotated IR 4. GlobalizeObjectGraph.cpp + tests - use new "private" attributes to create "private" IR. - also, tweak some of the op deleting mechanics, which was triggering some memory errors / assertions With this, we can run the classifier through and inline it as follows: ``` frontends/pytorch/utils/pt_util.py --import --exported-name forward ~/tmp/classifier.pt \ | npcomp-opt -torch-globalize-object-graph -inline ``` IR: https://gist.github.com/silvasean/32dcad9f6270557f412094a77cecdd69
2021-02-20 08:21:21 +08:00
IValueImporter(MlirBlock importBlock, MlirContext context,
ClassAnnotator &annotator)
: importBlock(importBlock), context(context), typeMapper(context),
annotator(annotator) {}
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
Get simple quantized model importing. This is enough to import the program and get it through the compilation pipeline. It of course fails at the VerifyBackendContract pass since there is a lot missing, but the final IR for a simple quantized MLP is looking pretty decent already: [IR](https://gist.github.com/silvasean/f76bccd76e9b193d396cfb2f9a11f54d) Main changes: - Add support for importing torch quantized tensors, including `torch.per_tensor_affine.create` op and `!torch.qint8` element type. - Add support for importing `LinearPackedParamsBase` (basically a weight + optional bias, but requires `torch.linear_params.create` op + `!torch.LinearParams` type to model it). This was less painful than I expected, as it has the necessary methods to opaquely unpack itself. I factored things so it should be easy to extend to other custom classes like `ConvPackedParamsBase`. - Add minimal boilerplate for importing `quantized::*` ops, with `quantized::linear` being a motivating example. - Add e2e test with simple quantized MLP (courtesy of @phoenix-meadowlark). This is somewhat of an abuse of `!numpy.ndarray` / `tensor`, as really the proper semantics of `!torch.qint8` dtype on a Torch tensor is "check the quantizer object of the tensor for side data (scale/offset, possibly per-channel) that defines the full semantics of the tensor". We don't have any such notion of "side data" for `!numpy.ndarray` / `tensor`, let alone anything that would have the associated behavior of keying off the dtype to determine if the side data is present. This will be fixed by a proper `!torch.tensor` type.
2021-05-20 02:40:48 +08:00
MlirValue importIValue(c10::IValue ivalue);
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
private:
Get simple quantized model importing. This is enough to import the program and get it through the compilation pipeline. It of course fails at the VerifyBackendContract pass since there is a lot missing, but the final IR for a simple quantized MLP is looking pretty decent already: [IR](https://gist.github.com/silvasean/f76bccd76e9b193d396cfb2f9a11f54d) Main changes: - Add support for importing torch quantized tensors, including `torch.per_tensor_affine.create` op and `!torch.qint8` element type. - Add support for importing `LinearPackedParamsBase` (basically a weight + optional bias, but requires `torch.linear_params.create` op + `!torch.LinearParams` type to model it). This was less painful than I expected, as it has the necessary methods to opaquely unpack itself. I factored things so it should be easy to extend to other custom classes like `ConvPackedParamsBase`. - Add minimal boilerplate for importing `quantized::*` ops, with `quantized::linear` being a motivating example. - Add e2e test with simple quantized MLP (courtesy of @phoenix-meadowlark). This is somewhat of an abuse of `!numpy.ndarray` / `tensor`, as really the proper semantics of `!torch.qint8` dtype on a Torch tensor is "check the quantizer object of the tensor for side data (scale/offset, possibly per-channel) that defines the full semantics of the tensor". We don't have any such notion of "side data" for `!numpy.ndarray` / `tensor`, let alone anything that would have the associated behavior of keying off the dtype to determine if the side data is present. This will be fixed by a proper `!torch.tensor` type.
2021-05-20 02:40:48 +08:00
MlirValue rawImportIValue(c10::IValue ivalue);
MlirValue importTensor(c10::IValue ivalue);
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
MlirValue importModule(torch::jit::Module jitModule);
Add ability to annotate TorchScript classes. The first use case is to annotate certain program constructs as either exported or private. In this commit we plumb it down to GlobalizeObjectGraph which makes use of this information. Recommended review order: 1. class_annotator.h/.cpp + `test/module_import/annotations/*` - New abstractions to communicate with Python code and annotate. 2. IR changes in TorchOps.td - Adding "private" attribute to various things. 3. ivalue_import.cpp changes - Module + ClassAnnotator = annotated IR 4. GlobalizeObjectGraph.cpp + tests - use new "private" attributes to create "private" IR. - also, tweak some of the op deleting mechanics, which was triggering some memory errors / assertions With this, we can run the classifier through and inline it as follows: ``` frontends/pytorch/utils/pt_util.py --import --exported-name forward ~/tmp/classifier.pt \ | npcomp-opt -torch-globalize-object-graph -inline ``` IR: https://gist.github.com/silvasean/32dcad9f6270557f412094a77cecdd69
2021-02-20 08:21:21 +08:00
void importMethod(torch::jit::Function *function, MlirBlock classTypeBody,
const MethodAnnotation &methodAnnotation);
Implement GlobalizeObjectGraph transformation. This required restructuring of how we model TorchScript on import. The main difference is that now we split out a `torch.class_type` that holds methods and declarations of the types of each slot. This is more consistent with TorchScript (our previous representation was "denormalized"). Recommended reading order: 1. check out the description of `torch.class_type` in `TorchOps.td` and look at `test/Dialect/Torch/ops.mlir` and `frontends/pytorch/test/module_import/` to familiarize with the new representation. - Just look at the new IR. The diff between the old names and new names is confusing. 2. check out `test/Dialect/Torch/globalize-object-graph*.mlir` and read along with the pass description in `include/npcomp/Dialect/Torch/Transforms/Passes.td` 3. Read the code in `GlobalizeObjectGraph.cpp` and miscellaneous changes in `ivalue_importer.cpp`, `TorchOps.cpp`, etc.
2021-02-18 03:28:51 +08:00
void importClassType(c10::ClassType *classType);
Properly import the entire torch::jit::CompilationUnit 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](https://github.com/pytorch/pytorch/blob/1d6bd157902d4b1347a5d03122d02b407658e263/torch/csrc/jit/runtime/interpreter.cpp#L937). We do import `prim::CallFunction` as a `std.call_indirect` though because it's more braindead to do it that way (it gets canonicalized to a direct call easily).
2021-02-27 08:20:35 +08:00
void importCompilationUnit(torch::jit::CompilationUnit *cu);
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
MlirBlock importBlock;
MlirContext context;
TypeMapper typeMapper;
Add ability to annotate TorchScript classes. The first use case is to annotate certain program constructs as either exported or private. In this commit we plumb it down to GlobalizeObjectGraph which makes use of this information. Recommended review order: 1. class_annotator.h/.cpp + `test/module_import/annotations/*` - New abstractions to communicate with Python code and annotate. 2. IR changes in TorchOps.td - Adding "private" attribute to various things. 3. ivalue_import.cpp changes - Module + ClassAnnotator = annotated IR 4. GlobalizeObjectGraph.cpp + tests - use new "private" attributes to create "private" IR. - also, tweak some of the op deleting mechanics, which was triggering some memory errors / assertions With this, we can run the classifier through and inline it as follows: ``` frontends/pytorch/utils/pt_util.py --import --exported-name forward ~/tmp/classifier.pt \ | npcomp-opt -torch-globalize-object-graph -inline ``` IR: https://gist.github.com/silvasean/32dcad9f6270557f412094a77cecdd69
2021-02-20 08:21:21 +08:00
ClassAnnotator &annotator;
Handle object identity correctly. This required some careful considerations when defining object identity for tensors. See the comments for how we do it. This also tracks some basic information for diagnostics.
2021-02-06 09:57:38 +08:00
// Map tracking already-imported values.
std::unordered_map<c10::IValue, MlirValue, IValueHasher, IValueEq> valueMap;
Properly import the entire torch::jit::CompilationUnit 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](https://github.com/pytorch/pytorch/blob/1d6bd157902d4b1347a5d03122d02b407658e263/torch/csrc/jit/runtime/interpreter.cpp#L937). We do import `prim::CallFunction` as a `std.call_indirect` though because it's more braindead to do it that way (it gets canonicalized to a direct call easily).
2021-02-27 08:20:35 +08:00
// The unique compilation unit that is shared by all modules reachable
// from the root ivalue being imported.
// It basically contains a symbol table of functions which are referenced from
// e.g. methods (the function names are meaningful and match with Python's
// module hierarchy, with the exception of `__main__` being replaced with
// `__torch__`).
torch::jit::CompilationUnit *compilationUnit = nullptr;
Handle object identity correctly. This required some careful considerations when defining object identity for tensors. See the comments for how we do it. This also tracks some basic information for diagnostics.
2021-02-06 09:57:38 +08:00
// Used to detect potentially aliasing tensors.
std::unordered_set<c10::StorageImpl *> seenStorageImpls;
Implement GlobalizeObjectGraph transformation. This required restructuring of how we model TorchScript on import. The main difference is that now we split out a `torch.class_type` that holds methods and declarations of the types of each slot. This is more consistent with TorchScript (our previous representation was "denormalized"). Recommended reading order: 1. check out the description of `torch.class_type` in `TorchOps.td` and look at `test/Dialect/Torch/ops.mlir` and `frontends/pytorch/test/module_import/` to familiarize with the new representation. - Just look at the new IR. The diff between the old names and new names is confusing. 2. check out `test/Dialect/Torch/globalize-object-graph*.mlir` and read along with the pass description in `include/npcomp/Dialect/Torch/Transforms/Passes.td` 3. Read the code in `GlobalizeObjectGraph.cpp` and miscellaneous changes in `ivalue_importer.cpp`, `TorchOps.cpp`, etc.
2021-02-18 03:28:51 +08:00
// The set of ClassType's that have already been imported.
//
// ClassType's are referenced via their `classType->name()->qualifiedName()`
// string (as an MLIR symbol name) so we don't need to keep a map associating
// them with the MlirOperation that they import into.
std::unordered_set<c10::ClassType *> classTypes;
Handle object identity correctly. This required some careful considerations when defining object identity for tensors. See the comments for how we do it. This also tracks some basic information for diagnostics.
2021-02-06 09:57:38 +08:00
// The stack of attribute names we have traversed to reach the current IValue.
// Used for diagnostics.
std::vector<std::string> attributeNameStack;
// The root module encountered during recursive IValue traversal.
// Used for diagnostics.
// Note that the "top-level" object being imported can in theory be a list
// of modules, so this is populated when our recursive traversal enters a
// module not enclosed in any other module, and unset after our recursive
// traversal exits the module.
c10::optional<std::string> rootModuleName;
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
};
} // namespace
Shore up error reporting for TorchScript import. This code was not exception safe -- it would leave an operation unattached to anything, which breaks MLIR's C++ data structure invariants (e.g. it cannot safely erase ops). Also, print out both the exception and any diagnostics, since they can both contain useful information.
2021-05-19 03:48:22 +08:00
MlirValue
IValueImporter::importModule(torch::jit::Module currentModule) {
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
// TODO: Can we do better?
MlirLocation loc = mlirLocationUnknownGet(context);
Implement GlobalizeObjectGraph transformation. This required restructuring of how we model TorchScript on import. The main difference is that now we split out a `torch.class_type` that holds methods and declarations of the types of each slot. This is more consistent with TorchScript (our previous representation was "denormalized"). Recommended reading order: 1. check out the description of `torch.class_type` in `TorchOps.td` and look at `test/Dialect/Torch/ops.mlir` and `frontends/pytorch/test/module_import/` to familiarize with the new representation. - Just look at the new IR. The diff between the old names and new names is confusing. 2. check out `test/Dialect/Torch/globalize-object-graph*.mlir` and read along with the pass description in `include/npcomp/Dialect/Torch/Transforms/Passes.td` 3. Read the code in `GlobalizeObjectGraph.cpp` and miscellaneous changes in `ivalue_importer.cpp`, `TorchOps.cpp`, etc.
2021-02-18 03:28:51 +08:00
c10::optional<c10::QualifiedName> maybeName = currentModule.type()->name();
if (!maybeName) {
throw std::invalid_argument("cannot import unnamed module");
}
std::string moduleTypeName = maybeName->qualifiedName();
Properly import the entire torch::jit::CompilationUnit 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](https://github.com/pytorch/pytorch/blob/1d6bd157902d4b1347a5d03122d02b407658e263/torch/csrc/jit/runtime/interpreter.cpp#L937). We do import `prim::CallFunction` as a `std.call_indirect` though because it's more braindead to do it that way (it gets canonicalized to a direct call easily).
2021-02-27 08:20:35 +08:00
// If this is the first time we are encountering a module, import the
// compilation unit.
importCompilationUnit(currentModule._ivalue()->compilation_unit().get());
Implement GlobalizeObjectGraph transformation. This required restructuring of how we model TorchScript on import. The main difference is that now we split out a `torch.class_type` that holds methods and declarations of the types of each slot. This is more consistent with TorchScript (our previous representation was "denormalized"). Recommended reading order: 1. check out the description of `torch.class_type` in `TorchOps.td` and look at `test/Dialect/Torch/ops.mlir` and `frontends/pytorch/test/module_import/` to familiarize with the new representation. - Just look at the new IR. The diff between the old names and new names is confusing. 2. check out `test/Dialect/Torch/globalize-object-graph*.mlir` and read along with the pass description in `include/npcomp/Dialect/Torch/Transforms/Passes.td` 3. Read the code in `GlobalizeObjectGraph.cpp` and miscellaneous changes in `ivalue_importer.cpp`, `TorchOps.cpp`, etc.
2021-02-18 03:28:51 +08:00
// Ensure the class type has been imported.
importClassType(currentModule.type().get());
MlirOperation nnModule = createMlirOperation(
"torch.nn_module", loc,
npcompNnModuleTypeGet(context, toMlirStringRef(moduleTypeName)),
mlirRegionCreate());
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
MlirRegion nnModuleRegion = mlirOperationGetRegion(nnModule, 0);
mlirRegionAppendOwnedBlock(nnModuleRegion, mlirBlockCreate(0, nullptr));
MlirBlock nnModuleBody = mlirRegionGetFirstBlock(nnModuleRegion);
Shore up error reporting for TorchScript import. This code was not exception safe -- it would leave an operation unattached to anything, which breaks MLIR's C++ data structure invariants (e.g. it cannot safely erase ops). Also, print out both the exception and any diagnostics, since they can both contain useful information.
2021-05-19 03:48:22 +08:00
auto inserter = caffe2::MakeGuard([&]() {
mlirBlockInsertOwnedOperationBefore(
importBlock, mlirBlockGetTerminator(importBlock), nnModule);
});
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
Handle object identity correctly. This required some careful considerations when defining object identity for tensors. See the comments for how we do it. This also tracks some basic information for diagnostics.
2021-02-06 09:57:38 +08:00
if (!rootModuleName.has_value()) {
Implement GlobalizeObjectGraph transformation. This required restructuring of how we model TorchScript on import. The main difference is that now we split out a `torch.class_type` that holds methods and declarations of the types of each slot. This is more consistent with TorchScript (our previous representation was "denormalized"). Recommended reading order: 1. check out the description of `torch.class_type` in `TorchOps.td` and look at `test/Dialect/Torch/ops.mlir` and `frontends/pytorch/test/module_import/` to familiarize with the new representation. - Just look at the new IR. The diff between the old names and new names is confusing. 2. check out `test/Dialect/Torch/globalize-object-graph*.mlir` and read along with the pass description in `include/npcomp/Dialect/Torch/Transforms/Passes.td` 3. Read the code in `GlobalizeObjectGraph.cpp` and miscellaneous changes in `ivalue_importer.cpp`, `TorchOps.cpp`, etc.
2021-02-18 03:28:51 +08:00
rootModuleName = moduleTypeName;
Handle object identity correctly. This required some careful considerations when defining object identity for tensors. See the comments for how we do it. This also tracks some basic information for diagnostics.
2021-02-06 09:57:38 +08:00
}
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
const std::vector<c10::IValue> &slots = currentModule._ivalue()->slots();
const std::vector<c10::ClassAttribute> &classAttributes =
currentModule.type()->getAttributes();
assert(slots.size() == classAttributes.size() &&
"mismatch between object and type!");
for (int i = 0, e = slots.size(); i < e; i++) {
const c10::ClassAttribute &classAttribute = classAttributes[i];
Handle object identity correctly. This required some careful considerations when defining object identity for tensors. See the comments for how we do it. This also tracks some basic information for diagnostics.
2021-02-06 09:57:38 +08:00
attributeNameStack.push_back(classAttribute.getName());
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
MlirValue slotValue = importIValue(slots[i]);
// TODO: Is it necessary to track whether an attribute is a "parameter"?
createMlirOperationAtEnd(
Implement GlobalizeObjectGraph transformation. This required restructuring of how we model TorchScript on import. The main difference is that now we split out a `torch.class_type` that holds methods and declarations of the types of each slot. This is more consistent with TorchScript (our previous representation was "denormalized"). Recommended reading order: 1. check out the description of `torch.class_type` in `TorchOps.td` and look at `test/Dialect/Torch/ops.mlir` and `frontends/pytorch/test/module_import/` to familiarize with the new representation. - Just look at the new IR. The diff between the old names and new names is confusing. 2. check out `test/Dialect/Torch/globalize-object-graph*.mlir` and read along with the pass description in `include/npcomp/Dialect/Torch/Transforms/Passes.td` 3. Read the code in `GlobalizeObjectGraph.cpp` and miscellaneous changes in `ivalue_importer.cpp`, `TorchOps.cpp`, etc.
2021-02-18 03:28:51 +08:00
nnModuleBody, "torch.slot", loc, slotValue,
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
toMlirNamedAttribute(
"name", mlirStringAttrGet(
context, toMlirStringRef(classAttribute.getName()))));
Handle object identity correctly. This required some careful considerations when defining object identity for tensors. See the comments for how we do it. This also tracks some basic information for diagnostics.
2021-02-06 09:57:38 +08:00
attributeNameStack.pop_back();
}
if (rootModuleName.has_value()) {
rootModuleName = c10::nullopt;
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
}
createMlirOperationAtEnd(nnModuleBody, "torch.nn_module_terminator", loc);
return mlirOperationGetResult(nnModule, 0);
}
MlirValue IValueImporter::importIValue(c10::IValue ivalue) {
Handle object identity correctly. This required some careful considerations when defining object identity for tensors. See the comments for how we do it. This also tracks some basic information for diagnostics.
2021-02-06 09:57:38 +08:00
auto it = valueMap.find(ivalue);
if (it != valueMap.end()) {
return it->second;
}
// Reject potentially aliased tensors.
if (ivalue.isTensor()) {
c10::StorageImpl *storageImpl =
ivalue.toTensor().storage().unsafeGetStorageImpl();
if (!seenStorageImpls.insert(storageImpl).second) {
std::stringstream msg;
msg << "Unhandled tensor that shares storage with another tensor.";
if (rootModuleName) {
msg << "\nFound at path '<root>."
<< c10::QualifiedName(attributeNameStack).qualifiedName()
<< "' from root object '" << *rootModuleName << "'";
}
throw std::invalid_argument(msg.str());
}
}
MlirValue value = rawImportIValue(ivalue);
valueMap[ivalue] = value;
return value;
}
MlirValue IValueImporter::rawImportIValue(c10::IValue ivalue) {
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
// TODO: Can we do better?
MlirLocation loc = mlirLocationUnknownGet(context);
if (ivalue.isBool()) {
MlirType type = npcompBoolTypeGet(context);
MlirOperation operation = createMlirOperationAtEnd(
importBlock, "basicpy.bool_constant", loc, type,
toMlirNamedAttribute("value",
mlirBoolAttrGet(context, ivalue.toBool())));
return mlirOperationGetResult(operation, 0);
}
if (ivalue.isDouble()) {
MlirType type = mlirF64TypeGet(context);
MlirOperation operation = createMlirOperationAtEnd(
importBlock, "basicpy.numeric_constant", loc, type,
toMlirNamedAttribute(
"value", mlirFloatAttrDoubleGet(context, type, ivalue.toDouble())));
return mlirOperationGetResult(operation, 0);
}
if (ivalue.isInt()) {
MlirType type = mlirIntegerTypeGet(context, 64);
MlirOperation operation = createMlirOperationAtEnd(
importBlock, "basicpy.numeric_constant", loc, type,
toMlirNamedAttribute("value",
mlirIntegerAttrGet(type, ivalue.toInt())));
return mlirOperationGetResult(operation, 0);
}
Handle object identity correctly. This required some careful considerations when defining object identity for tensors. See the comments for how we do it. This also tracks some basic information for diagnostics.
2021-02-06 09:57:38 +08:00
if (ivalue.isList()) {
c10::List<c10::IValue> list = ivalue.toList();
std::vector<MlirValue> elems;
for (const c10::IValue &elem : list) {
elems.push_back(importIValue(elem));
}
MlirOperation operation =
Add TorchList type and prim::ListConstruct #218
2021-06-05 06:57:21 +08:00
createMlirOperationAtEnd(importBlock, "torch.prim.ListConstruct", loc,
npcompListTypeGet(
typeMapper.mapFromTorchType(
loc, list.elementType())), elems);
Handle object identity correctly. This required some careful considerations when defining object identity for tensors. See the comments for how we do it. This also tracks some basic information for diagnostics.
2021-02-06 09:57:38 +08:00
return mlirOperationGetResult(operation, 0);
}
Add Conv2D Torchscript Import Support (#167) Adds support for lowering a torch.nn.Conv2d module to the Torch Dialect through TorchScript import. Generated IR can be viewed here: https://gist.github.com/brycearden/6c0f790115c4577249372ef82768e6fd Required implementing support for tuple in the ivalue importer and list in the node importer.
2021-02-26 04:14:00 +08:00
if (ivalue.isTuple()) {
auto list = ivalue.toTuple()->elements();
std::vector<MlirValue> elems;
for (const c10::IValue &elem : list) {
elems.push_back(importIValue(elem));
}
MlirOperation operation =
createMlirOperationAtEnd(importBlock, "basicpy.build_tuple", loc,
npcompTupleTypeGet(context), elems);
return mlirOperationGetResult(operation, 0);
}
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
if (ivalue.isTensor()) {
Get simple quantized model importing. This is enough to import the program and get it through the compilation pipeline. It of course fails at the VerifyBackendContract pass since there is a lot missing, but the final IR for a simple quantized MLP is looking pretty decent already: [IR](https://gist.github.com/silvasean/f76bccd76e9b193d396cfb2f9a11f54d) Main changes: - Add support for importing torch quantized tensors, including `torch.per_tensor_affine.create` op and `!torch.qint8` element type. - Add support for importing `LinearPackedParamsBase` (basically a weight + optional bias, but requires `torch.linear_params.create` op + `!torch.LinearParams` type to model it). This was less painful than I expected, as it has the necessary methods to opaquely unpack itself. I factored things so it should be easy to extend to other custom classes like `ConvPackedParamsBase`. - Add minimal boilerplate for importing `quantized::*` ops, with `quantized::linear` being a motivating example. - Add e2e test with simple quantized MLP (courtesy of @phoenix-meadowlark). This is somewhat of an abuse of `!numpy.ndarray` / `tensor`, as really the proper semantics of `!torch.qint8` dtype on a Torch tensor is "check the quantizer object of the tensor for side data (scale/offset, possibly per-channel) that defines the full semantics of the tensor". We don't have any such notion of "side data" for `!numpy.ndarray` / `tensor`, let alone anything that would have the associated behavior of keying off the dtype to determine if the side data is present. This will be fixed by a proper `!torch.tensor` type.
2021-05-20 02:40:48 +08:00
return importTensor(ivalue);
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
}
if (ivalue.isModule()) {
return importModule(ivalue.toModule());
}
[torch] Add import support for IValue string Type(s) (#179) * [torch] Add import support for IValue string Type(s) * [test] Add test for Strings import
2021-03-05 05:08:50 +08:00
if (ivalue.isString()) {
MlirType type = npcompBytesTypeGet(context);
MlirOperation operation = createMlirOperationAtEnd(
importBlock, "basicpy.bytes_constant", loc, type,
toMlirNamedAttribute(
"value",
mlirStringAttrGet(context,
toMlirStringRef(ivalue.toString()->string()))));
return mlirOperationGetResult(operation, 0);
}
Add NoneType support for ivalue_importer PyTorch added a Global variable `_is_full_backward_hook` recently. See https://github.com/pytorch/pytorch/pull/46163 Signed-off-by: Bairen Yi <yibairen.byron@bytedance.com>
2021-02-18 21:56:15 +08:00
if (ivalue.isNone()) {
Introduce native `!torch.none` type. - Add `torch.constant.none` op to construct it (naming is chosen to be analogous to Torch's representation of a prim::Constant with NoneType, rather than using the "singleton" terminology of Basicpy).
2021-06-15 02:36:10 +08:00
MlirOperation operation =
createMlirOperationAtEnd(importBlock, "torch.constant.none", loc,
npcompTorchNoneTypeGet(context));
Add NoneType support for ivalue_importer PyTorch added a Global variable `_is_full_backward_hook` recently. See https://github.com/pytorch/pytorch/pull/46163 Signed-off-by: Bairen Yi <yibairen.byron@bytedance.com>
2021-02-18 21:56:15 +08:00
return mlirOperationGetResult(operation, 0);
}
Get simple quantized model importing. This is enough to import the program and get it through the compilation pipeline. It of course fails at the VerifyBackendContract pass since there is a lot missing, but the final IR for a simple quantized MLP is looking pretty decent already: [IR](https://gist.github.com/silvasean/f76bccd76e9b193d396cfb2f9a11f54d) Main changes: - Add support for importing torch quantized tensors, including `torch.per_tensor_affine.create` op and `!torch.qint8` element type. - Add support for importing `LinearPackedParamsBase` (basically a weight + optional bias, but requires `torch.linear_params.create` op + `!torch.LinearParams` type to model it). This was less painful than I expected, as it has the necessary methods to opaquely unpack itself. I factored things so it should be easy to extend to other custom classes like `ConvPackedParamsBase`. - Add minimal boilerplate for importing `quantized::*` ops, with `quantized::linear` being a motivating example. - Add e2e test with simple quantized MLP (courtesy of @phoenix-meadowlark). This is somewhat of an abuse of `!numpy.ndarray` / `tensor`, as really the proper semantics of `!torch.qint8` dtype on a Torch tensor is "check the quantizer object of the tensor for side data (scale/offset, possibly per-channel) that defines the full semantics of the tensor". We don't have any such notion of "side data" for `!numpy.ndarray` / `tensor`, let alone anything that would have the associated behavior of keying off the dtype to determine if the side data is present. This will be fixed by a proper `!torch.tensor` type.
2021-05-20 02:40:48 +08:00
if (ivalue.isCustomClass()) {
if (ivalue.type().get() ==
c10::getCustomClassType<c10::intrusive_ptr<LinearPackedParamsBase>>()
.get()) {
c10::intrusive_ptr<LinearPackedParamsBase> linearParams =
ivalue.toCustomClass<LinearPackedParamsBase>();
at::Tensor weight;
c10::optional<at::Tensor> bias;
std::tie(weight, bias) = linearParams->unpack();
MlirValue weightValue = importIValue(c10::IValue(weight));
c10::optional<MlirValue> biasValue = c10::nullopt;
if (bias.has_value()) {
biasValue = importIValue(c10::IValue(*bias));
}
MlirOperation operation = createMlirOperationAtEnd(
importBlock, "torch.linear_params.create", loc,
npcompLinearParamsTypeGet(context), weightValue, biasValue);
return mlirOperationGetResult(operation, 0);
}
}
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
std::stringstream msg;
msg << "Unsupported ivalue: " << ivalue;
throw std::invalid_argument(msg.str());
}
Get simple quantized model importing. This is enough to import the program and get it through the compilation pipeline. It of course fails at the VerifyBackendContract pass since there is a lot missing, but the final IR for a simple quantized MLP is looking pretty decent already: [IR](https://gist.github.com/silvasean/f76bccd76e9b193d396cfb2f9a11f54d) Main changes: - Add support for importing torch quantized tensors, including `torch.per_tensor_affine.create` op and `!torch.qint8` element type. - Add support for importing `LinearPackedParamsBase` (basically a weight + optional bias, but requires `torch.linear_params.create` op + `!torch.LinearParams` type to model it). This was less painful than I expected, as it has the necessary methods to opaquely unpack itself. I factored things so it should be easy to extend to other custom classes like `ConvPackedParamsBase`. - Add minimal boilerplate for importing `quantized::*` ops, with `quantized::linear` being a motivating example. - Add e2e test with simple quantized MLP (courtesy of @phoenix-meadowlark). This is somewhat of an abuse of `!numpy.ndarray` / `tensor`, as really the proper semantics of `!torch.qint8` dtype on a Torch tensor is "check the quantizer object of the tensor for side data (scale/offset, possibly per-channel) that defines the full semantics of the tensor". We don't have any such notion of "side data" for `!numpy.ndarray` / `tensor`, let alone anything that would have the associated behavior of keying off the dtype to determine if the side data is present. This will be fixed by a proper `!torch.tensor` type.
2021-05-20 02:40:48 +08:00
MlirValue IValueImporter::importTensor(c10::IValue ivalue) {
assert(ivalue.isTensor() && "expected a tensor!");
// TODO: Can we do better?
MlirLocation loc = mlirLocationUnknownGet(context);
// Import the bulk tensor representation.
at::Tensor tensor = ivalue.toTensor().contiguous();
Introduce `!torch.tensor` / `!torch.vtensor` types. This removes our reliance on the numpy dialect and avoids our off-label use of the builtin tnesor type for modeling unknown dtypes. The `!torch.vtensor` (`ValueTensorType`) type is a value-semantic tensor. The `!torch.tensor` (`NonValueTensorType`) type is a non-value-semantic tensor. The new types look as follows syntactically: ``` // Least-static-information, non-value-semantic tensor. !torch.tensor // Explicit form of least-static-information variant. !torch.tensor<*,unk> // Least-static-information, value-semantic tensor. !torch.vtensor // Explicit form of least-static-information variant. !torch.vtensor<*,unk> // Fixed-set of allowable element types, with first-class support for // Torch's frontend signedness semantics. !torch.tensor<*,si32> // First-class support for unknown dtypes. !torch.tensor<[?,?,?],unk> // Standard MLIR representation of `?` for unknown dimensions. !torch.tensor<[?,2,?,4],unk> // Statically shaped / dtyped example. !torch.vtensor<[1,2,3,4],f32> ``` This required fairly significant changes throughout the compiler, but overall it is a big cleanup. We now have a much clearer layering of "the Torch frontend lowering" vs "lowering to std + linalg + etc.". At the C++ level, there is `ValueTensorType`, `NonValueTensorType`. We also have a helper `BaseTensorType` (kind of like ShapedType) which interoperates with those two. Included changes: - New `torch.tensor(dense<0.0> : tensor<5xf32>) : !torch.tensor` op for creating torch tensor literals in the frontend. - Consistently use signedness for the types (except i1 which I didn't touch -- we need to sort out the situation with !basicpy.BoolType there anyway so will be attending to that soon) - Frontend can annotate whether an argument to the function has value semantics. We currently require this, as our backend contract does not currently allow us to even model the non-value-semantic case. Before, the value-semantic assumption was randomly injected in the middle of the pass pipeline. - Move ArrayToTensor (now called MaximizeValueSemantics) and RefinePublicReturn passes to torch dialect. - The TorchToStd and TorchToLinalg passes are now type conversions from `!torch.vtensor` to `tensor` and use the dialect conversion infra. The overall conversion pipeline is set up following the best practices of the "Type Conversions the Not-So-Hard Way" talk. This required introducing `torch-func-builtin-tensorize` and `torch-finalizing-builtin-tensorize` passes analogous to the upstream bufferization passes with the corresponding names (mostly just copypasta from there). - Misc Torch-level canonicalizations -- we now cleanly layer the lowering to std later in the pipeline, so we are gradually lessening our reliance on random std constant folding before we get to that point. Recommended review order: - New types in TorchTypes.td/TorchTypes.h/TorchDialect.cpp - New ops in TorchOps.td / TorchOps.cpp - Less important / more mechanical stuff - Frontend changes. - Pass changes/additions in `Torch/Transforms` and `Conversion/`
2021-05-21 08:07:18 +08:00
MlirAttribute denseElements = convertTensorToMlirElementsAttr(tensor, loc);
MlirOperation tensorOp =
createMlirOperationAtEnd(importBlock, "torch.tensor", loc,
npcompNonValueTensorTypeGetFromShaped(
mlirAttributeGetType(denseElements)),
toMlirNamedAttribute("value", denseElements));
MlirValue tensorReprValue = mlirOperationGetResult(tensorOp, 0);
Get simple quantized model importing. This is enough to import the program and get it through the compilation pipeline. It of course fails at the VerifyBackendContract pass since there is a lot missing, but the final IR for a simple quantized MLP is looking pretty decent already: [IR](https://gist.github.com/silvasean/f76bccd76e9b193d396cfb2f9a11f54d) Main changes: - Add support for importing torch quantized tensors, including `torch.per_tensor_affine.create` op and `!torch.qint8` element type. - Add support for importing `LinearPackedParamsBase` (basically a weight + optional bias, but requires `torch.linear_params.create` op + `!torch.LinearParams` type to model it). This was less painful than I expected, as it has the necessary methods to opaquely unpack itself. I factored things so it should be easy to extend to other custom classes like `ConvPackedParamsBase`. - Add minimal boilerplate for importing `quantized::*` ops, with `quantized::linear` being a motivating example. - Add e2e test with simple quantized MLP (courtesy of @phoenix-meadowlark). This is somewhat of an abuse of `!numpy.ndarray` / `tensor`, as really the proper semantics of `!torch.qint8` dtype on a Torch tensor is "check the quantizer object of the tensor for side data (scale/offset, possibly per-channel) that defines the full semantics of the tensor". We don't have any such notion of "side data" for `!numpy.ndarray` / `tensor`, let alone anything that would have the associated behavior of keying off the dtype to determine if the side data is present. This will be fixed by a proper `!torch.tensor` type.
2021-05-20 02:40:48 +08:00
// Construct the complete tensor value. This is trivial for most tensors, but
// for quantized tensors (and probably sparse too, TBD) there is more for us
// to do.
MlirValue tensorValue;
if (tensor.is_quantized()) {
// Note that Torch models quantization in a type-erased way. So we don't
// make an effort here to do any special static modeling. If desired, later
// compiler stages that are building a statically modeled quantization
// representation will need to convert this to their representation.
std::vector<int64_t> shape(tensor.sizes().begin(), tensor.sizes().end());
Introduce `!torch.tensor` / `!torch.vtensor` types. This removes our reliance on the numpy dialect and avoids our off-label use of the builtin tnesor type for modeling unknown dtypes. The `!torch.vtensor` (`ValueTensorType`) type is a value-semantic tensor. The `!torch.tensor` (`NonValueTensorType`) type is a non-value-semantic tensor. The new types look as follows syntactically: ``` // Least-static-information, non-value-semantic tensor. !torch.tensor // Explicit form of least-static-information variant. !torch.tensor<*,unk> // Least-static-information, value-semantic tensor. !torch.vtensor // Explicit form of least-static-information variant. !torch.vtensor<*,unk> // Fixed-set of allowable element types, with first-class support for // Torch's frontend signedness semantics. !torch.tensor<*,si32> // First-class support for unknown dtypes. !torch.tensor<[?,?,?],unk> // Standard MLIR representation of `?` for unknown dimensions. !torch.tensor<[?,2,?,4],unk> // Statically shaped / dtyped example. !torch.vtensor<[1,2,3,4],f32> ``` This required fairly significant changes throughout the compiler, but overall it is a big cleanup. We now have a much clearer layering of "the Torch frontend lowering" vs "lowering to std + linalg + etc.". At the C++ level, there is `ValueTensorType`, `NonValueTensorType`. We also have a helper `BaseTensorType` (kind of like ShapedType) which interoperates with those two. Included changes: - New `torch.tensor(dense<0.0> : tensor<5xf32>) : !torch.tensor` op for creating torch tensor literals in the frontend. - Consistently use signedness for the types (except i1 which I didn't touch -- we need to sort out the situation with !basicpy.BoolType there anyway so will be attending to that soon) - Frontend can annotate whether an argument to the function has value semantics. We currently require this, as our backend contract does not currently allow us to even model the non-value-semantic case. Before, the value-semantic assumption was randomly injected in the middle of the pass pipeline. - Move ArrayToTensor (now called MaximizeValueSemantics) and RefinePublicReturn passes to torch dialect. - The TorchToStd and TorchToLinalg passes are now type conversions from `!torch.vtensor` to `tensor` and use the dialect conversion infra. The overall conversion pipeline is set up following the best practices of the "Type Conversions the Not-So-Hard Way" talk. This required introducing `torch-func-builtin-tensorize` and `torch-finalizing-builtin-tensorize` passes analogous to the upstream bufferization passes with the corresponding names (mostly just copypasta from there). - Misc Torch-level canonicalizations -- we now cleanly layer the lowering to std later in the pipeline, so we are gradually lessening our reliance on random std constant folding before we get to that point. Recommended review order: - New types in TorchTypes.td/TorchTypes.h/TorchDialect.cpp - New ops in TorchOps.td / TorchOps.cpp - Less important / more mechanical stuff - Frontend changes. - Pass changes/additions in `Torch/Transforms` and `Conversion/`
2021-05-21 08:07:18 +08:00
MlirType quantizedTensorType = npcompNonValueTensorTypeGet(
context, shape.size(), shape.data(),
typeMapper.mapFromTorchScalarType(tensor.scalar_type()));
Get simple quantized model importing. This is enough to import the program and get it through the compilation pipeline. It of course fails at the VerifyBackendContract pass since there is a lot missing, but the final IR for a simple quantized MLP is looking pretty decent already: [IR](https://gist.github.com/silvasean/f76bccd76e9b193d396cfb2f9a11f54d) Main changes: - Add support for importing torch quantized tensors, including `torch.per_tensor_affine.create` op and `!torch.qint8` element type. - Add support for importing `LinearPackedParamsBase` (basically a weight + optional bias, but requires `torch.linear_params.create` op + `!torch.LinearParams` type to model it). This was less painful than I expected, as it has the necessary methods to opaquely unpack itself. I factored things so it should be easy to extend to other custom classes like `ConvPackedParamsBase`. - Add minimal boilerplate for importing `quantized::*` ops, with `quantized::linear` being a motivating example. - Add e2e test with simple quantized MLP (courtesy of @phoenix-meadowlark). This is somewhat of an abuse of `!numpy.ndarray` / `tensor`, as really the proper semantics of `!torch.qint8` dtype on a Torch tensor is "check the quantizer object of the tensor for side data (scale/offset, possibly per-channel) that defines the full semantics of the tensor". We don't have any such notion of "side data" for `!numpy.ndarray` / `tensor`, let alone anything that would have the associated behavior of keying off the dtype to determine if the side data is present. This will be fixed by a proper `!torch.tensor` type.
2021-05-20 02:40:48 +08:00
if (tensor.qscheme() == c10::kPerTensorAffine) {
MlirValue qScale = importIValue(c10::IValue(tensor.q_scale()));
MlirValue zeroPoint = importIValue(c10::IValue(tensor.q_zero_point()));
MlirOperation quantizedTensor = createMlirOperationAtEnd(
importBlock, "torch.per_tensor_affine.create", loc,
quantizedTensorType, tensorReprValue, qScale, zeroPoint);
tensorValue = mlirOperationGetResult(quantizedTensor, 0);
} else {
std::stringstream msg;
msg << "Unsupported quantization scheme '"
<< c10::toString(tensor.qscheme()) << "' for tensor: " << ivalue;
throw std::invalid_argument(msg.str());
}
} else {
tensorValue = tensorReprValue;
}
Introduce `!torch.tensor` / `!torch.vtensor` types. This removes our reliance on the numpy dialect and avoids our off-label use of the builtin tnesor type for modeling unknown dtypes. The `!torch.vtensor` (`ValueTensorType`) type is a value-semantic tensor. The `!torch.tensor` (`NonValueTensorType`) type is a non-value-semantic tensor. The new types look as follows syntactically: ``` // Least-static-information, non-value-semantic tensor. !torch.tensor // Explicit form of least-static-information variant. !torch.tensor<*,unk> // Least-static-information, value-semantic tensor. !torch.vtensor // Explicit form of least-static-information variant. !torch.vtensor<*,unk> // Fixed-set of allowable element types, with first-class support for // Torch's frontend signedness semantics. !torch.tensor<*,si32> // First-class support for unknown dtypes. !torch.tensor<[?,?,?],unk> // Standard MLIR representation of `?` for unknown dimensions. !torch.tensor<[?,2,?,4],unk> // Statically shaped / dtyped example. !torch.vtensor<[1,2,3,4],f32> ``` This required fairly significant changes throughout the compiler, but overall it is a big cleanup. We now have a much clearer layering of "the Torch frontend lowering" vs "lowering to std + linalg + etc.". At the C++ level, there is `ValueTensorType`, `NonValueTensorType`. We also have a helper `BaseTensorType` (kind of like ShapedType) which interoperates with those two. Included changes: - New `torch.tensor(dense<0.0> : tensor<5xf32>) : !torch.tensor` op for creating torch tensor literals in the frontend. - Consistently use signedness for the types (except i1 which I didn't touch -- we need to sort out the situation with !basicpy.BoolType there anyway so will be attending to that soon) - Frontend can annotate whether an argument to the function has value semantics. We currently require this, as our backend contract does not currently allow us to even model the non-value-semantic case. Before, the value-semantic assumption was randomly injected in the middle of the pass pipeline. - Move ArrayToTensor (now called MaximizeValueSemantics) and RefinePublicReturn passes to torch dialect. - The TorchToStd and TorchToLinalg passes are now type conversions from `!torch.vtensor` to `tensor` and use the dialect conversion infra. The overall conversion pipeline is set up following the best practices of the "Type Conversions the Not-So-Hard Way" talk. This required introducing `torch-func-builtin-tensorize` and `torch-finalizing-builtin-tensorize` passes analogous to the upstream bufferization passes with the corresponding names (mostly just copypasta from there). - Misc Torch-level canonicalizations -- we now cleanly layer the lowering to std later in the pipeline, so we are gradually lessening our reliance on random std constant folding before we get to that point. Recommended review order: - New types in TorchTypes.td/TorchTypes.h/TorchDialect.cpp - New ops in TorchOps.td / TorchOps.cpp - Less important / more mechanical stuff - Frontend changes. - Pass changes/additions in `Torch/Transforms` and `Conversion/`
2021-05-21 08:07:18 +08:00
return tensorValue;
Get simple quantized model importing. This is enough to import the program and get it through the compilation pipeline. It of course fails at the VerifyBackendContract pass since there is a lot missing, but the final IR for a simple quantized MLP is looking pretty decent already: [IR](https://gist.github.com/silvasean/f76bccd76e9b193d396cfb2f9a11f54d) Main changes: - Add support for importing torch quantized tensors, including `torch.per_tensor_affine.create` op and `!torch.qint8` element type. - Add support for importing `LinearPackedParamsBase` (basically a weight + optional bias, but requires `torch.linear_params.create` op + `!torch.LinearParams` type to model it). This was less painful than I expected, as it has the necessary methods to opaquely unpack itself. I factored things so it should be easy to extend to other custom classes like `ConvPackedParamsBase`. - Add minimal boilerplate for importing `quantized::*` ops, with `quantized::linear` being a motivating example. - Add e2e test with simple quantized MLP (courtesy of @phoenix-meadowlark). This is somewhat of an abuse of `!numpy.ndarray` / `tensor`, as really the proper semantics of `!torch.qint8` dtype on a Torch tensor is "check the quantizer object of the tensor for side data (scale/offset, possibly per-channel) that defines the full semantics of the tensor". We don't have any such notion of "side data" for `!numpy.ndarray` / `tensor`, let alone anything that would have the associated behavior of keying off the dtype to determine if the side data is present. This will be fixed by a proper `!torch.tensor` type.
2021-05-20 02:40:48 +08:00
}
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
void IValueImporter::importMethod(torch::jit::Function *function,
Add ability to annotate TorchScript classes. The first use case is to annotate certain program constructs as either exported or private. In this commit we plumb it down to GlobalizeObjectGraph which makes use of this information. Recommended review order: 1. class_annotator.h/.cpp + `test/module_import/annotations/*` - New abstractions to communicate with Python code and annotate. 2. IR changes in TorchOps.td - Adding "private" attribute to various things. 3. ivalue_import.cpp changes - Module + ClassAnnotator = annotated IR 4. GlobalizeObjectGraph.cpp + tests - use new "private" attributes to create "private" IR. - also, tweak some of the op deleting mechanics, which was triggering some memory errors / assertions With this, we can run the classifier through and inline it as follows: ``` frontends/pytorch/utils/pt_util.py --import --exported-name forward ~/tmp/classifier.pt \ | npcomp-opt -torch-globalize-object-graph -inline ``` IR: https://gist.github.com/silvasean/32dcad9f6270557f412094a77cecdd69
2021-02-20 08:21:21 +08:00
MlirBlock classTypeBody,
const MethodAnnotation &methodAnnotation) {
Properly import the entire torch::jit::CompilationUnit 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](https://github.com/pytorch/pytorch/blob/1d6bd157902d4b1347a5d03122d02b407658e263/torch/csrc/jit/runtime/interpreter.cpp#L937). We do import `prim::CallFunction` as a `std.call_indirect` though because it's more braindead to do it that way (it gets canonicalized to a direct call easily).
2021-02-27 08:20:35 +08:00
// The function's name becomes the MLIR symbol table name of the imported func
// when we import the compilation unit.
const std::string &symName = function->qualname().qualifiedName();
MlirAttribute functionSymbolRef =
mlirFlatSymbolRefAttrGet(context, toMlirStringRef(symName));
Add ability to annotate TorchScript classes. The first use case is to annotate certain program constructs as either exported or private. In this commit we plumb it down to GlobalizeObjectGraph which makes use of this information. Recommended review order: 1. class_annotator.h/.cpp + `test/module_import/annotations/*` - New abstractions to communicate with Python code and annotate. 2. IR changes in TorchOps.td - Adding "private" attribute to various things. 3. ivalue_import.cpp changes - Module + ClassAnnotator = annotated IR 4. GlobalizeObjectGraph.cpp + tests - use new "private" attributes to create "private" IR. - also, tweak some of the op deleting mechanics, which was triggering some memory errors / assertions With this, we can run the classifier through and inline it as follows: ``` frontends/pytorch/utils/pt_util.py --import --exported-name forward ~/tmp/classifier.pt \ | npcomp-opt -torch-globalize-object-graph -inline ``` IR: https://gist.github.com/silvasean/32dcad9f6270557f412094a77cecdd69
2021-02-20 08:21:21 +08:00
c10::optional<MlirNamedAttribute> isPrivate;
if (!methodAnnotation.isExported) {
isPrivate = toMlirNamedAttribute("isPrivate", mlirUnitAttrGet(context));
}
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
createMlirOperationAtEnd(
Implement GlobalizeObjectGraph transformation. This required restructuring of how we model TorchScript on import. The main difference is that now we split out a `torch.class_type` that holds methods and declarations of the types of each slot. This is more consistent with TorchScript (our previous representation was "denormalized"). Recommended reading order: 1. check out the description of `torch.class_type` in `TorchOps.td` and look at `test/Dialect/Torch/ops.mlir` and `frontends/pytorch/test/module_import/` to familiarize with the new representation. - Just look at the new IR. The diff between the old names and new names is confusing. 2. check out `test/Dialect/Torch/globalize-object-graph*.mlir` and read along with the pass description in `include/npcomp/Dialect/Torch/Transforms/Passes.td` 3. Read the code in `GlobalizeObjectGraph.cpp` and miscellaneous changes in `ivalue_importer.cpp`, `TorchOps.cpp`, etc.
2021-02-18 03:28:51 +08:00
classTypeBody, "torch.method", mlirLocationUnknownGet(context),
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
toMlirNamedAttribute(
"name",
mlirStringAttrGet(context, toMlirStringRef(function->name()))),
Properly import the entire torch::jit::CompilationUnit 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](https://github.com/pytorch/pytorch/blob/1d6bd157902d4b1347a5d03122d02b407658e263/torch/csrc/jit/runtime/interpreter.cpp#L937). We do import `prim::CallFunction` as a `std.call_indirect` though because it's more braindead to do it that way (it gets canonicalized to a direct call easily).
2021-02-27 08:20:35 +08:00
toMlirNamedAttribute("function", functionSymbolRef), isPrivate);
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
}
Implement GlobalizeObjectGraph transformation. This required restructuring of how we model TorchScript on import. The main difference is that now we split out a `torch.class_type` that holds methods and declarations of the types of each slot. This is more consistent with TorchScript (our previous representation was "denormalized"). Recommended reading order: 1. check out the description of `torch.class_type` in `TorchOps.td` and look at `test/Dialect/Torch/ops.mlir` and `frontends/pytorch/test/module_import/` to familiarize with the new representation. - Just look at the new IR. The diff between the old names and new names is confusing. 2. check out `test/Dialect/Torch/globalize-object-graph*.mlir` and read along with the pass description in `include/npcomp/Dialect/Torch/Transforms/Passes.td` 3. Read the code in `GlobalizeObjectGraph.cpp` and miscellaneous changes in `ivalue_importer.cpp`, `TorchOps.cpp`, etc.
2021-02-18 03:28:51 +08:00
void IValueImporter::importClassType(c10::ClassType *classType) {
if (!classTypes.insert(classType).second) {
return;
}
// TODO: Can we do better?
MlirLocation loc = mlirLocationUnknownGet(context);
MlirOperation op = createMlirOperationAtEnd(
importBlock, "torch.class_type", loc, mlirRegionCreate(),
toMlirNamedAttribute(
"sym_name",
mlirStringAttrGet(
context, toMlirStringRef(classType->name()->qualifiedName()))));
MlirRegion region = mlirOperationGetRegion(op, 0);
mlirRegionAppendOwnedBlock(region, mlirBlockCreate(0, nullptr));
MlirBlock classTypeBody = mlirRegionGetFirstBlock(region);
Add ability to annotate TorchScript classes. The first use case is to annotate certain program constructs as either exported or private. In this commit we plumb it down to GlobalizeObjectGraph which makes use of this information. Recommended review order: 1. class_annotator.h/.cpp + `test/module_import/annotations/*` - New abstractions to communicate with Python code and annotate. 2. IR changes in TorchOps.td - Adding "private" attribute to various things. 3. ivalue_import.cpp changes - Module + ClassAnnotator = annotated IR 4. GlobalizeObjectGraph.cpp + tests - use new "private" attributes to create "private" IR. - also, tweak some of the op deleting mechanics, which was triggering some memory errors / assertions With this, we can run the classifier through and inline it as follows: ``` frontends/pytorch/utils/pt_util.py --import --exported-name forward ~/tmp/classifier.pt \ | npcomp-opt -torch-globalize-object-graph -inline ``` IR: https://gist.github.com/silvasean/32dcad9f6270557f412094a77cecdd69
2021-02-20 08:21:21 +08:00
ClassAnnotation &classAnnotation =
annotator.getOrCreateClassAnnotation(classType);
const auto &attributeAnnotations = classAnnotation.getAttributeAnnotations();
const auto &classAttributes = classType->getAttributes();
for (int i = 0, e = classAttributes.size(); i != e; i++) {
const c10::ClassAttribute &classAttribute = classAttributes[i];
c10::optional<MlirNamedAttribute> isPrivate;
if (!attributeAnnotations[i].isExported) {
isPrivate = toMlirNamedAttribute("isPrivate", mlirUnitAttrGet(context));
}
Implement GlobalizeObjectGraph transformation. This required restructuring of how we model TorchScript on import. The main difference is that now we split out a `torch.class_type` that holds methods and declarations of the types of each slot. This is more consistent with TorchScript (our previous representation was "denormalized"). Recommended reading order: 1. check out the description of `torch.class_type` in `TorchOps.td` and look at `test/Dialect/Torch/ops.mlir` and `frontends/pytorch/test/module_import/` to familiarize with the new representation. - Just look at the new IR. The diff between the old names and new names is confusing. 2. check out `test/Dialect/Torch/globalize-object-graph*.mlir` and read along with the pass description in `include/npcomp/Dialect/Torch/Transforms/Passes.td` 3. Read the code in `GlobalizeObjectGraph.cpp` and miscellaneous changes in `ivalue_importer.cpp`, `TorchOps.cpp`, etc.
2021-02-18 03:28:51 +08:00
createMlirOperationAtEnd(
classTypeBody, "torch.attr", loc,
toMlirNamedAttribute(
"name", mlirStringAttrGet(
context, toMlirStringRef(classAttribute.getName()))),
toMlirNamedAttribute("type",
mlirTypeAttrGet(typeMapper.mapFromTorchType(
Add ability to annotate TorchScript classes. The first use case is to annotate certain program constructs as either exported or private. In this commit we plumb it down to GlobalizeObjectGraph which makes use of this information. Recommended review order: 1. class_annotator.h/.cpp + `test/module_import/annotations/*` - New abstractions to communicate with Python code and annotate. 2. IR changes in TorchOps.td - Adding "private" attribute to various things. 3. ivalue_import.cpp changes - Module + ClassAnnotator = annotated IR 4. GlobalizeObjectGraph.cpp + tests - use new "private" attributes to create "private" IR. - also, tweak some of the op deleting mechanics, which was triggering some memory errors / assertions With this, we can run the classifier through and inline it as follows: ``` frontends/pytorch/utils/pt_util.py --import --exported-name forward ~/tmp/classifier.pt \ | npcomp-opt -torch-globalize-object-graph -inline ``` IR: https://gist.github.com/silvasean/32dcad9f6270557f412094a77cecdd69
2021-02-20 08:21:21 +08:00
loc, classAttribute.getType()))),
isPrivate);
Implement GlobalizeObjectGraph transformation. This required restructuring of how we model TorchScript on import. The main difference is that now we split out a `torch.class_type` that holds methods and declarations of the types of each slot. This is more consistent with TorchScript (our previous representation was "denormalized"). Recommended reading order: 1. check out the description of `torch.class_type` in `TorchOps.td` and look at `test/Dialect/Torch/ops.mlir` and `frontends/pytorch/test/module_import/` to familiarize with the new representation. - Just look at the new IR. The diff between the old names and new names is confusing. 2. check out `test/Dialect/Torch/globalize-object-graph*.mlir` and read along with the pass description in `include/npcomp/Dialect/Torch/Transforms/Passes.td` 3. Read the code in `GlobalizeObjectGraph.cpp` and miscellaneous changes in `ivalue_importer.cpp`, `TorchOps.cpp`, etc.
2021-02-18 03:28:51 +08:00
}
Add ability to annotate TorchScript classes. The first use case is to annotate certain program constructs as either exported or private. In this commit we plumb it down to GlobalizeObjectGraph which makes use of this information. Recommended review order: 1. class_annotator.h/.cpp + `test/module_import/annotations/*` - New abstractions to communicate with Python code and annotate. 2. IR changes in TorchOps.td - Adding "private" attribute to various things. 3. ivalue_import.cpp changes - Module + ClassAnnotator = annotated IR 4. GlobalizeObjectGraph.cpp + tests - use new "private" attributes to create "private" IR. - also, tweak some of the op deleting mechanics, which was triggering some memory errors / assertions With this, we can run the classifier through and inline it as follows: ``` frontends/pytorch/utils/pt_util.py --import --exported-name forward ~/tmp/classifier.pt \ | npcomp-opt -torch-globalize-object-graph -inline ``` IR: https://gist.github.com/silvasean/32dcad9f6270557f412094a77cecdd69
2021-02-20 08:21:21 +08:00
const auto &methodAnnotations = classAnnotation.getMethodAnnotations();
const auto &methods = classType->methods();
for (int i = 0, e = methods.size(); i != e; i++) {
importMethod(methods[i], classTypeBody, methodAnnotations[i]);
Implement GlobalizeObjectGraph transformation. This required restructuring of how we model TorchScript on import. The main difference is that now we split out a `torch.class_type` that holds methods and declarations of the types of each slot. This is more consistent with TorchScript (our previous representation was "denormalized"). Recommended reading order: 1. check out the description of `torch.class_type` in `TorchOps.td` and look at `test/Dialect/Torch/ops.mlir` and `frontends/pytorch/test/module_import/` to familiarize with the new representation. - Just look at the new IR. The diff between the old names and new names is confusing. 2. check out `test/Dialect/Torch/globalize-object-graph*.mlir` and read along with the pass description in `include/npcomp/Dialect/Torch/Transforms/Passes.td` 3. Read the code in `GlobalizeObjectGraph.cpp` and miscellaneous changes in `ivalue_importer.cpp`, `TorchOps.cpp`, etc.
2021-02-18 03:28:51 +08:00
}
createMlirOperationAtEnd(classTypeBody, "torch.class_type_terminator", loc);
}
Properly import the entire torch::jit::CompilationUnit 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](https://github.com/pytorch/pytorch/blob/1d6bd157902d4b1347a5d03122d02b407658e263/torch/csrc/jit/runtime/interpreter.cpp#L937). We do import `prim::CallFunction` as a `std.call_indirect` though because it's more braindead to do it that way (it gets canonicalized to a direct call easily).
2021-02-27 08:20:35 +08:00
void IValueImporter::importCompilationUnit(torch::jit::CompilationUnit *cu) {
if (compilationUnit == nullptr) {
compilationUnit = cu;
} else {
// All sorts of stuff is connected to the compilation unit, such as
// c10::ClassType's (owned by the compilation unit), c10::FunctionType
// (which holds a pointer to a torch::jit::Function in the compilation
// unit), load-bearing symbol table names of functions, etc.
//
// It doesn't seem to be defined how multiple compilation units semantically
// connect with each other, and it doesn't seem to happen either (though
// structurally at the C++ level nothing prevents it), so make it an error.
if (compilationUnit != cu) {
throw std::invalid_argument(
"found two compilation units while importing");
}
return;
}
for (torch::jit::Function *function : cu->get_functions()) {
Miscellaneous changes while trying to work on ResNet18 - Move frontend lowering pipelines to c++ (this helps with reproducing failures in npcomp-opt) - Add debugging printouts when compilation fails on RefBackendTestConfig The experience now when a test fails during MLIR lowering is now like this: ``` NPCOMP TorchScript Object Graph IR -> NPCOMP Backend IR lowering failed with the following diagnostics: failed to legalize operation 'torch.global_slot' Module does not conform to npcomp's backend contract. See dialect conversion legality information above. Error can be reproduced with: $ npcomp-opt -torchscript-to-npcomp-backend-pipeline /tmp/ResNet18Module.mlir ``` And when TorchScript->MLIR import fails it looks like this: ``` PyTorch TorchScript module -> NPCOMP Object Graph IR import failed with the following diagnostics: unhandled prim operation: %18 : int = prim::min(%17) # /usr/local/google/home/silvasean/.local/lib/python3.9/site-packages/torch/nn/functional.py:4532:4 ``` Also, - Add `--filter=<regex>` to e2e test harness to filter tests. - Add a few prim ops that were needed to import ResNet18 - Fix torch.prim.Loop.condition assemblyFormat (it previously would not round-trip in the case of no loop-carried variables)
2021-04-22 06:07:15 +08:00
// Useful for debugging errors in free functions that end up being
// unused. These can be missing when round-tripping through the on-disk
// format, even though they still cause import issues when importing
// through the larger Python session where they originate.
// std::cerr << "NAME: " << function->qualname().qualifiedName() << "\n";
// std::cerr << *function->graph();
Basic infra for annotate shapes and dtypes on arguments. These allow users to annotate a known "type bound" on the argument, which can seed shape/dtype inference. We don't rewrite the function types as part of the import process (it will happen in a yet-to-be-written pass) because: 1. We would need to interprocedurally rewrite all calls to keep the IR consistent. Currently, we have a place after GlobalizeObjectGraph but before we convert to tensors where this is convenient to do. Ideally, we would do this on the object graph representation. 1. We don't necessarily know that adjusting the function type is a legal calling convention change. The pass will have blessed knowledge (by the pass pipeline author) that adjusting the argument type based on the type bound is safe (which it frequently is). 2. Note that in principle, a type bound could be a fairly general thing (such as maximum sizes of dimensions, unions of multiple concrete types, etc.). The pass will in principle have logic to interpret the type bounds and to determine a suitable "best" (and legal) argument type.
2021-03-31 07:11:41 +08:00
MethodAnnotation *annotation =
annotator.getMethodAnnotationForFunction(function);
MlirOperation func = importJitFunctionAsFuncOp(
context, function, [&](int argIndex) -> MlirAttribute {
if (!annotation || !annotation->argAnnotations.has_value()) {
return {nullptr};
}
Introduce `!torch.tensor` / `!torch.vtensor` types. This removes our reliance on the numpy dialect and avoids our off-label use of the builtin tnesor type for modeling unknown dtypes. The `!torch.vtensor` (`ValueTensorType`) type is a value-semantic tensor. The `!torch.tensor` (`NonValueTensorType`) type is a non-value-semantic tensor. The new types look as follows syntactically: ``` // Least-static-information, non-value-semantic tensor. !torch.tensor // Explicit form of least-static-information variant. !torch.tensor<*,unk> // Least-static-information, value-semantic tensor. !torch.vtensor // Explicit form of least-static-information variant. !torch.vtensor<*,unk> // Fixed-set of allowable element types, with first-class support for // Torch's frontend signedness semantics. !torch.tensor<*,si32> // First-class support for unknown dtypes. !torch.tensor<[?,?,?],unk> // Standard MLIR representation of `?` for unknown dimensions. !torch.tensor<[?,2,?,4],unk> // Statically shaped / dtyped example. !torch.vtensor<[1,2,3,4],f32> ``` This required fairly significant changes throughout the compiler, but overall it is a big cleanup. We now have a much clearer layering of "the Torch frontend lowering" vs "lowering to std + linalg + etc.". At the C++ level, there is `ValueTensorType`, `NonValueTensorType`. We also have a helper `BaseTensorType` (kind of like ShapedType) which interoperates with those two. Included changes: - New `torch.tensor(dense<0.0> : tensor<5xf32>) : !torch.tensor` op for creating torch tensor literals in the frontend. - Consistently use signedness for the types (except i1 which I didn't touch -- we need to sort out the situation with !basicpy.BoolType there anyway so will be attending to that soon) - Frontend can annotate whether an argument to the function has value semantics. We currently require this, as our backend contract does not currently allow us to even model the non-value-semantic case. Before, the value-semantic assumption was randomly injected in the middle of the pass pipeline. - Move ArrayToTensor (now called MaximizeValueSemantics) and RefinePublicReturn passes to torch dialect. - The TorchToStd and TorchToLinalg passes are now type conversions from `!torch.vtensor` to `tensor` and use the dialect conversion infra. The overall conversion pipeline is set up following the best practices of the "Type Conversions the Not-So-Hard Way" talk. This required introducing `torch-func-builtin-tensorize` and `torch-finalizing-builtin-tensorize` passes analogous to the upstream bufferization passes with the corresponding names (mostly just copypasta from there). - Misc Torch-level canonicalizations -- we now cleanly layer the lowering to std later in the pipeline, so we are gradually lessening our reliance on random std constant folding before we get to that point. Recommended review order: - New types in TorchTypes.td/TorchTypes.h/TorchDialect.cpp - New ops in TorchOps.td / TorchOps.cpp - Less important / more mechanical stuff - Frontend changes. - Pass changes/additions in `Torch/Transforms` and `Conversion/`
2021-05-21 08:07:18 +08:00
c10::optional<std::vector<int64_t>> &maybeShape =
annotation->argAnnotations.value()[argIndex].shape;
c10::optional<c10::ScalarType> &maybeDtype =
annotation->argAnnotations.value()[argIndex].dtype;
bool hasValueSemantics =
annotation->argAnnotations.value()[argIndex].hasValueSemantics;
Basic infra for annotate shapes and dtypes on arguments. These allow users to annotate a known "type bound" on the argument, which can seed shape/dtype inference. We don't rewrite the function types as part of the import process (it will happen in a yet-to-be-written pass) because: 1. We would need to interprocedurally rewrite all calls to keep the IR consistent. Currently, we have a place after GlobalizeObjectGraph but before we convert to tensors where this is convenient to do. Ideally, we would do this on the object graph representation. 1. We don't necessarily know that adjusting the function type is a legal calling convention change. The pass will have blessed knowledge (by the pass pipeline author) that adjusting the argument type based on the type bound is safe (which it frequently is). 2. Note that in principle, a type bound could be a fairly general thing (such as maximum sizes of dimensions, unions of multiple concrete types, etc.). The pass will in principle have logic to interpret the type bounds and to determine a suitable "best" (and legal) argument type.
2021-03-31 07:11:41 +08:00
// TODO: Handle unranked tensors and tensors with unknown dtype (but
// possibly known ranks/sizes).
Introduce `!torch.tensor` / `!torch.vtensor` types. This removes our reliance on the numpy dialect and avoids our off-label use of the builtin tnesor type for modeling unknown dtypes. The `!torch.vtensor` (`ValueTensorType`) type is a value-semantic tensor. The `!torch.tensor` (`NonValueTensorType`) type is a non-value-semantic tensor. The new types look as follows syntactically: ``` // Least-static-information, non-value-semantic tensor. !torch.tensor // Explicit form of least-static-information variant. !torch.tensor<*,unk> // Least-static-information, value-semantic tensor. !torch.vtensor // Explicit form of least-static-information variant. !torch.vtensor<*,unk> // Fixed-set of allowable element types, with first-class support for // Torch's frontend signedness semantics. !torch.tensor<*,si32> // First-class support for unknown dtypes. !torch.tensor<[?,?,?],unk> // Standard MLIR representation of `?` for unknown dimensions. !torch.tensor<[?,2,?,4],unk> // Statically shaped / dtyped example. !torch.vtensor<[1,2,3,4],f32> ``` This required fairly significant changes throughout the compiler, but overall it is a big cleanup. We now have a much clearer layering of "the Torch frontend lowering" vs "lowering to std + linalg + etc.". At the C++ level, there is `ValueTensorType`, `NonValueTensorType`. We also have a helper `BaseTensorType` (kind of like ShapedType) which interoperates with those two. Included changes: - New `torch.tensor(dense<0.0> : tensor<5xf32>) : !torch.tensor` op for creating torch tensor literals in the frontend. - Consistently use signedness for the types (except i1 which I didn't touch -- we need to sort out the situation with !basicpy.BoolType there anyway so will be attending to that soon) - Frontend can annotate whether an argument to the function has value semantics. We currently require this, as our backend contract does not currently allow us to even model the non-value-semantic case. Before, the value-semantic assumption was randomly injected in the middle of the pass pipeline. - Move ArrayToTensor (now called MaximizeValueSemantics) and RefinePublicReturn passes to torch dialect. - The TorchToStd and TorchToLinalg passes are now type conversions from `!torch.vtensor` to `tensor` and use the dialect conversion infra. The overall conversion pipeline is set up following the best practices of the "Type Conversions the Not-So-Hard Way" talk. This required introducing `torch-func-builtin-tensorize` and `torch-finalizing-builtin-tensorize` passes analogous to the upstream bufferization passes with the corresponding names (mostly just copypasta from there). - Misc Torch-level canonicalizations -- we now cleanly layer the lowering to std later in the pipeline, so we are gradually lessening our reliance on random std constant folding before we get to that point. Recommended review order: - New types in TorchTypes.td/TorchTypes.h/TorchDialect.cpp - New ops in TorchOps.td / TorchOps.cpp - Less important / more mechanical stuff - Frontend changes. - Pass changes/additions in `Torch/Transforms` and `Conversion/`
2021-05-21 08:07:18 +08:00
if (!maybeShape || !maybeDtype) {
Basic infra for annotate shapes and dtypes on arguments. These allow users to annotate a known "type bound" on the argument, which can seed shape/dtype inference. We don't rewrite the function types as part of the import process (it will happen in a yet-to-be-written pass) because: 1. We would need to interprocedurally rewrite all calls to keep the IR consistent. Currently, we have a place after GlobalizeObjectGraph but before we convert to tensors where this is convenient to do. Ideally, we would do this on the object graph representation. 1. We don't necessarily know that adjusting the function type is a legal calling convention change. The pass will have blessed knowledge (by the pass pipeline author) that adjusting the argument type based on the type bound is safe (which it frequently is). 2. Note that in principle, a type bound could be a fairly general thing (such as maximum sizes of dimensions, unions of multiple concrete types, etc.). The pass will in principle have logic to interpret the type bounds and to determine a suitable "best" (and legal) argument type.
2021-03-31 07:11:41 +08:00
return {nullptr};
}
Introduce `!torch.tensor` / `!torch.vtensor` types. This removes our reliance on the numpy dialect and avoids our off-label use of the builtin tnesor type for modeling unknown dtypes. The `!torch.vtensor` (`ValueTensorType`) type is a value-semantic tensor. The `!torch.tensor` (`NonValueTensorType`) type is a non-value-semantic tensor. The new types look as follows syntactically: ``` // Least-static-information, non-value-semantic tensor. !torch.tensor // Explicit form of least-static-information variant. !torch.tensor<*,unk> // Least-static-information, value-semantic tensor. !torch.vtensor // Explicit form of least-static-information variant. !torch.vtensor<*,unk> // Fixed-set of allowable element types, with first-class support for // Torch's frontend signedness semantics. !torch.tensor<*,si32> // First-class support for unknown dtypes. !torch.tensor<[?,?,?],unk> // Standard MLIR representation of `?` for unknown dimensions. !torch.tensor<[?,2,?,4],unk> // Statically shaped / dtyped example. !torch.vtensor<[1,2,3,4],f32> ``` This required fairly significant changes throughout the compiler, but overall it is a big cleanup. We now have a much clearer layering of "the Torch frontend lowering" vs "lowering to std + linalg + etc.". At the C++ level, there is `ValueTensorType`, `NonValueTensorType`. We also have a helper `BaseTensorType` (kind of like ShapedType) which interoperates with those two. Included changes: - New `torch.tensor(dense<0.0> : tensor<5xf32>) : !torch.tensor` op for creating torch tensor literals in the frontend. - Consistently use signedness for the types (except i1 which I didn't touch -- we need to sort out the situation with !basicpy.BoolType there anyway so will be attending to that soon) - Frontend can annotate whether an argument to the function has value semantics. We currently require this, as our backend contract does not currently allow us to even model the non-value-semantic case. Before, the value-semantic assumption was randomly injected in the middle of the pass pipeline. - Move ArrayToTensor (now called MaximizeValueSemantics) and RefinePublicReturn passes to torch dialect. - The TorchToStd and TorchToLinalg passes are now type conversions from `!torch.vtensor` to `tensor` and use the dialect conversion infra. The overall conversion pipeline is set up following the best practices of the "Type Conversions the Not-So-Hard Way" talk. This required introducing `torch-func-builtin-tensorize` and `torch-finalizing-builtin-tensorize` passes analogous to the upstream bufferization passes with the corresponding names (mostly just copypasta from there). - Misc Torch-level canonicalizations -- we now cleanly layer the lowering to std later in the pipeline, so we are gradually lessening our reliance on random std constant folding before we get to that point. Recommended review order: - New types in TorchTypes.td/TorchTypes.h/TorchDialect.cpp - New ops in TorchOps.td / TorchOps.cpp - Less important / more mechanical stuff - Frontend changes. - Pass changes/additions in `Torch/Transforms` and `Conversion/`
2021-05-21 08:07:18 +08:00
std::vector<int64_t> shape = *maybeShape;
MlirType dtype = TypeMapper(context).mapFromTorchScalarType(
mlirLocationUnknownGet(context), *maybeDtype);
MlirType typeBound;
if (hasValueSemantics) {
typeBound = npcompValueTensorTypeGet(context, shape.size(),
shape.data(), dtype);
} else {
typeBound = npcompNonValueTensorTypeGet(context, shape.size(),
shape.data(), dtype);
}
Basic infra for annotate shapes and dtypes on arguments. These allow users to annotate a known "type bound" on the argument, which can seed shape/dtype inference. We don't rewrite the function types as part of the import process (it will happen in a yet-to-be-written pass) because: 1. We would need to interprocedurally rewrite all calls to keep the IR consistent. Currently, we have a place after GlobalizeObjectGraph but before we convert to tensors where this is convenient to do. Ideally, we would do this on the object graph representation. 1. We don't necessarily know that adjusting the function type is a legal calling convention change. The pass will have blessed knowledge (by the pass pipeline author) that adjusting the argument type based on the type bound is safe (which it frequently is). 2. Note that in principle, a type bound could be a fairly general thing (such as maximum sizes of dimensions, unions of multiple concrete types, etc.). The pass will in principle have logic to interpret the type bounds and to determine a suitable "best" (and legal) argument type.
2021-03-31 07:11:41 +08:00
MlirNamedAttribute typeBoundAttr = toMlirNamedAttribute(
"torch.type_bound", mlirTypeAttrGet(typeBound));
return mlirDictionaryAttrGet(context, 1, &typeBoundAttr);
});
Properly import the entire torch::jit::CompilationUnit 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](https://github.com/pytorch/pytorch/blob/1d6bd157902d4b1347a5d03122d02b407658e263/torch/csrc/jit/runtime/interpreter.cpp#L937). We do import `prim::CallFunction` as a `std.call_indirect` though because it's more braindead to do it that way (it gets canonicalized to a direct call easily).
2021-02-27 08:20:35 +08:00
// For IValue importing, the logical linkage structure of the module
// is determined by the object graph.
//
// The functions' symbol names are thus irrelevant to the module's
// externally visible characteristics, so mark them all as private.
//
// These functions may be referenced by the object graph, which can make
// them reachable from the exernally visible characteristics of the module,
// but they cannot be intrinsically externally visible.
mlirOperationSetAttributeByName(
func, toMlirStringRef("sym_visibility"),
mlirStringAttrGet(context, toMlirStringRef("private")));
mlirBlockInsertOwnedOperationBefore(
importBlock, mlirBlockGetTerminator(importBlock), func);
}
}
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
void torch_mlir::importIValue(c10::IValue ivalue, MlirBlock block,
Add ability to annotate TorchScript classes. The first use case is to annotate certain program constructs as either exported or private. In this commit we plumb it down to GlobalizeObjectGraph which makes use of this information. Recommended review order: 1. class_annotator.h/.cpp + `test/module_import/annotations/*` - New abstractions to communicate with Python code and annotate. 2. IR changes in TorchOps.td - Adding "private" attribute to various things. 3. ivalue_import.cpp changes - Module + ClassAnnotator = annotated IR 4. GlobalizeObjectGraph.cpp + tests - use new "private" attributes to create "private" IR. - also, tweak some of the op deleting mechanics, which was triggering some memory errors / assertions With this, we can run the classifier through and inline it as follows: ``` frontends/pytorch/utils/pt_util.py --import --exported-name forward ~/tmp/classifier.pt \ | npcomp-opt -torch-globalize-object-graph -inline ``` IR: https://gist.github.com/silvasean/32dcad9f6270557f412094a77cecdd69
2021-02-20 08:21:21 +08:00
MlirContext context, ClassAnnotator &annotator) {
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
// When debugging module importing, it can be useful to dump as so:
// if (ivalue.isModule())
Extend GlobalizeObjectGraph to handle torch.prim.GetAttr returning NnModuleType This happens in practice. With this, we can globalize slots for the non-trivial classifier layer obtained from https://github.com/NVIDIA/NeMo/blob/main/tutorials/nlp/Joint_Intent_and_Slot_Classification.ipynb This also adds support for tuple return types, which were needed by that model.
2021-02-19 09:10:17 +08:00
// ivalue.toModule().dump(true, false, false);
Add ability to annotate TorchScript classes. The first use case is to annotate certain program constructs as either exported or private. In this commit we plumb it down to GlobalizeObjectGraph which makes use of this information. Recommended review order: 1. class_annotator.h/.cpp + `test/module_import/annotations/*` - New abstractions to communicate with Python code and annotate. 2. IR changes in TorchOps.td - Adding "private" attribute to various things. 3. ivalue_import.cpp changes - Module + ClassAnnotator = annotated IR 4. GlobalizeObjectGraph.cpp + tests - use new "private" attributes to create "private" IR. - also, tweak some of the op deleting mechanics, which was triggering some memory errors / assertions With this, we can run the classifier through and inline it as follows: ``` frontends/pytorch/utils/pt_util.py --import --exported-name forward ~/tmp/classifier.pt \ | npcomp-opt -torch-globalize-object-graph -inline ``` IR: https://gist.github.com/silvasean/32dcad9f6270557f412094a77cecdd69
2021-02-20 08:21:21 +08:00
IValueImporter importer(block, context, annotator);
Add initial TorchScript module importer It turns out that this was easiest to structure as a general IValue importer, since torch module are just one of the possible IValue's. We import the IValue object graph in a braindead fashion into basicpy ops and a new `torch.nn_module` op that is used to model the attributes/methods of a torch::jit::Module IValue. See `Torch/ops.mlir` for an example, and also check out the .py import tests in `frontends/pytorch/test/module_import`. As part of this change, a few housekeeping tasks: - extract some helpers from graph_importer.cpp - more helpers around the C API - misc touchups
2021-01-28 08:35:44 +08:00
importer.importIValue(ivalue);
}