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
|
|
|
|
|
2020-12-12 06:43:38 +08:00
|
|
|
#include "mlir/IR/BuiltinTypes.h"
|
2020-09-29 03:02:35 +08:00
|
|
|
#include "mlir/IR/OpDefinition.h"
|
|
|
|
#include "mlir/IR/OpImplementation.h"
|
2021-01-28 08:35:44 +08:00
|
|
|
#include "mlir/IR/SymbolTable.h"
|
2021-03-02 07:00:32 +08:00
|
|
|
#include "mlir/Interfaces/ControlFlowInterfaces.h"
|
2021-03-02 09:24:15 +08:00
|
|
|
#include "mlir/Interfaces/SideEffectInterfaces.h"
|
2020-10-23 14:31:34 +08:00
|
|
|
#include "npcomp/Dialect/Torch/IR/OpInterfaces.h"
|
2021-01-28 08:35:44 +08:00
|
|
|
#include "npcomp/Dialect/Torch/IR/TorchTypes.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; }
|
|
|
|
};
|
|
|
|
|
2020-09-29 03:02:35 +08:00
|
|
|
#endif // NPCOMP_DIALECT_TORCH_IR_TORCHOPS_H
|