torch-mlir/include/npcomp/Dialect/Torch/IR/TorchOps.h

91 lines
3.4 KiB
C
Raw Normal View History

2020-09-29 03:02:35 +08:00
//===------------------------------------------------------------*- C++ -*-===//
//
// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
#ifndef NPCOMP_DIALECT_TORCH_IR_TORCHOPS_H
#define NPCOMP_DIALECT_TORCH_IR_TORCHOPS_H
#include "mlir/IR/BuiltinTypes.h"
2020-09-29 03:02:35 +08:00
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/OpImplementation.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/Interfaces/CastInterfaces.h"
#include "mlir/Interfaces/ControlFlowInterfaces.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 "mlir/Interfaces/SideEffectInterfaces.h"
#include "npcomp/Dialect/Torch/IR/OpInterfaces.h"
#include "npcomp/Dialect/Torch/IR/TorchTypes.h"
#include "npcomp/Interfaces/Traits.h"
2020-09-29 03:02:35 +08:00
#define GET_OP_CLASSES
#include "npcomp/Dialect/Torch/IR/TorchOps.h.inc"
Support multiple instances of a class in GlobalizeObjectGraph. This happens in practice with e.g. ResNet from torchvision (multiple instances of the same BatchNorm class). The key observation is that for this program, and the expected set of programs, we can convert the program to the same globalized form with a bit more static analysis and effort to suitably monomorphize the program. Though what we are doing here is fairly annoying to implement, it saves any nontrivial later pass from having to do similar analyses (or worse). E.g. shape inference would need to be object-graph aware, mutation/lifetime analyses would have to be aware, etc. Additionally, it would make us front-load what it means to have a !torch.nn.Module type on an ABI boundary, which we are just not ready to handle. I'm really, really hoping that in practice we can get away with this, otherwise it's going to be really rough designing a representation (and implementing everything to back it) that is convenient to transform and gracefully scales from full object graph (in the most dynamic case) down to a fixed set of global slots like we have here (in the most static case, which we presume a lot of practical programs fall into). This also involved introducing a `torch-prepare-for-globalize-object-graph` pass that does a minimal set of lowerings to simplify the IR into a more orthogonal and analyzable form, and a `torch-globalize-pipeline` helper. Recommended review order: - updated documentation in Passes.td - new tests in `globalize-object-graph-multiple-instances*.mlir` - implementation of GlobalizeObjectGraph.cpp - PrepareForGlobalizeObjectGraph.cpp + prepare-for-globalize-object-graph.mlir - misc stuff like torch-globalize-pipeline pipeline definition. With this, we can import, globalize, and inline resnet18 from torchvision: https://gist.github.com/silvasean/821586afc19b67d9fb72030b2e0adeb8
2021-03-10 12:33:21 +08:00
template <> struct llvm::DenseMapInfo<::mlir::NPCOMP::Torch::SlotOp> {
using SlotOp = ::mlir::NPCOMP::Torch::SlotOp;
static SlotOp getEmptyKey() {
auto *pointer = llvm::DenseMapInfo<void *>::getEmptyKey();
return SlotOp::getFromOpaquePointer(pointer);
}
static SlotOp getTombstoneKey() {
auto *pointer = llvm::DenseMapInfo<void *>::getTombstoneKey();
return SlotOp::getFromOpaquePointer(pointer);
}
static unsigned getHashValue(SlotOp val) {
return hash_value(val.getAsOpaquePointer());
}
static bool isEqual(SlotOp lhs, SlotOp rhs) { return lhs == rhs; }
};
template <> struct llvm::DenseMapInfo<::mlir::NPCOMP::Torch::NnModuleOp> {
using NnModuleOp = ::mlir::NPCOMP::Torch::NnModuleOp;
static NnModuleOp getEmptyKey() {
auto *pointer = llvm::DenseMapInfo<void *>::getEmptyKey();
return NnModuleOp::getFromOpaquePointer(pointer);
}
static NnModuleOp getTombstoneKey() {
auto *pointer = llvm::DenseMapInfo<void *>::getTombstoneKey();
return NnModuleOp::getFromOpaquePointer(pointer);
}
static unsigned getHashValue(NnModuleOp val) {
return hash_value(val.getAsOpaquePointer());
}
static bool isEqual(NnModuleOp lhs, NnModuleOp rhs) { return lhs == rhs; }
};
template <> struct llvm::DenseMapInfo<::mlir::NPCOMP::Torch::ClassTypeOp> {
using ClassTypeOp = ::mlir::NPCOMP::Torch::ClassTypeOp;
static ClassTypeOp getEmptyKey() {
auto *pointer = llvm::DenseMapInfo<void *>::getEmptyKey();
return ClassTypeOp::getFromOpaquePointer(pointer);
}
static ClassTypeOp getTombstoneKey() {
auto *pointer = llvm::DenseMapInfo<void *>::getTombstoneKey();
return ClassTypeOp::getFromOpaquePointer(pointer);
}
static unsigned getHashValue(ClassTypeOp val) {
return hash_value(val.getAsOpaquePointer());
}
static bool isEqual(ClassTypeOp lhs, ClassTypeOp rhs) { return lhs == rhs; }
};
template <> struct llvm::DenseMapInfo<::mlir::NPCOMP::Torch::GlobalSlotOp> {
using OpTy = ::mlir::NPCOMP::Torch::GlobalSlotOp;
static OpTy getEmptyKey() {
auto *pointer = llvm::DenseMapInfo<void *>::getEmptyKey();
return OpTy::getFromOpaquePointer(pointer);
}
static OpTy getTombstoneKey() {
auto *pointer = llvm::DenseMapInfo<void *>::getTombstoneKey();
return OpTy::getFromOpaquePointer(pointer);
}
static unsigned getHashValue(OpTy val) {
return hash_value(val.getAsOpaquePointer());
}
static bool isEqual(OpTy lhs, OpTy rhs) { return lhs == rhs; }
};
2020-09-29 03:02:35 +08:00
#endif // NPCOMP_DIALECT_TORCH_IR_TORCHOPS_H