torch-mlir/lib/InitAll.cpp

56 lines
2.2 KiB
C++
Raw Normal View History

//===----------------------------------------------------------------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
#include "npcomp/InitAll.h"
#include "mlir/IR/Dialect.h"
Add support for compiling through IREE. Recommended review order: - Changes in frontends/pytorch/examples/ - Changes in python/npcomp/compiler/pytorch/backend/ - Boilerplate for the `npcomp-iree-backend-lower-linkage` pass. This change separates out a `npcomp.compiler.pytorch.backend.frontend_lowering` module that does the common lowering for all backends. The individual compiler backends `npcomp.compiler.pytorch.backend.{refjit,iree}` now accept a loosely defined "TCP + scalar code" IR mix that will be formalized in the future as the interface to codegen backends. This also required adding a small pass `npcomp-iree-backend-lower-linkage` which adds `iree.module.export` onto functions, and layering that into the frontend flow. The pass doesn't require a C++-level dependency on IREE, which is nice for now. TBD how we are going to handle lists (we hope we can get away with sneakerneting some td files and relying on loose IR compatibility). Running through IREE requires the ability to import `iree.compiler` and `iree.runtime`, which can be obtained as follows: ``` python3 -m pip install iree-compiler-snapshot iree-runtime-snapshot -f https://github.com/google/iree/releases/tag/snapshot-20210406.200 PYTHONPATH="${PYTHONPATH}:${MY_IREE_BUILD}/bindings/python/" ``` This patch makes it painfully clear that we don't have any e2e testing harness to really plug into, and also don't have a usable Python API to our compiler stack (something usable in a jupyter notebook). That will be addressed in subsequent commits. We've been flying by the seat of our pants with this `examples` directory that isn't subject to any kind of testing or real usability concerns.
2021-04-09 04:05:16 +08:00
#include "npcomp/Backend/IREE/Passes.h"
#include "npcomp/Conversion/Passes.h"
#include "npcomp/Dialect/ATen/IR/ATenDialect.h"
#include "npcomp/Dialect/ATen/Transforms/Passes.h"
#include "npcomp/Dialect/Basicpy/IR/BasicpyDialect.h"
#include "npcomp/Dialect/Basicpy/Transforms/Passes.h"
#include "npcomp/Dialect/Numpy/IR/NumpyDialect.h"
#include "npcomp/Dialect/Numpy/Transforms/Passes.h"
#include "npcomp/Dialect/Refback/IR/RefbackDialect.h"
#include "npcomp/Dialect/Refbackrt/IR/RefbackrtDialect.h"
#include "npcomp/Dialect/TCF/IR/TCFDialect.h"
#include "npcomp/Dialect/TCF/Transforms/Passes.h"
#include "npcomp/Dialect/TCP/IR/TCPDialect.h"
#include "npcomp/Dialect/TCP/Transforms/Passes.h"
2020-09-29 03:02:35 +08:00
#include "npcomp/Dialect/Torch/IR/TorchDialect.h"
#include "npcomp/Dialect/Torch/Transforms/Passes.h"
#include "npcomp/RefBackend/RefBackend.h"
#include "npcomp/Typing/Transforms/Passes.h"
void mlir::NPCOMP::registerAllDialects(mlir::DialectRegistry &registry) {
// clang-format off
registry.insert<mlir::NPCOMP::aten::ATenDialect,
Basicpy::BasicpyDialect,
Numpy::NumpyDialect,
refbackrt::RefbackrtDialect,
refback::RefbackDialect,
tcf::TCFDialect,
2020-09-29 03:02:35 +08:00
tcp::TCPDialect,
mlir::NPCOMP::Torch::TorchDialect>();
// clang-format on
}
void mlir::NPCOMP::registerAllPasses() {
mlir::NPCOMP::registerATenPasses();
mlir::NPCOMP::registerRefBackendPasses();
mlir::NPCOMP::registerConversionPasses();
mlir::NPCOMP::registerBasicpyPasses();
mlir::NPCOMP::registerNumpyPasses();
mlir::NPCOMP::registerTCFPasses();
mlir::NPCOMP::registerTCPPasses();
mlir::NPCOMP::registerTorchPasses();
mlir::NPCOMP::registerTypingPasses();
Add support for compiling through IREE. Recommended review order: - Changes in frontends/pytorch/examples/ - Changes in python/npcomp/compiler/pytorch/backend/ - Boilerplate for the `npcomp-iree-backend-lower-linkage` pass. This change separates out a `npcomp.compiler.pytorch.backend.frontend_lowering` module that does the common lowering for all backends. The individual compiler backends `npcomp.compiler.pytorch.backend.{refjit,iree}` now accept a loosely defined "TCP + scalar code" IR mix that will be formalized in the future as the interface to codegen backends. This also required adding a small pass `npcomp-iree-backend-lower-linkage` which adds `iree.module.export` onto functions, and layering that into the frontend flow. The pass doesn't require a C++-level dependency on IREE, which is nice for now. TBD how we are going to handle lists (we hope we can get away with sneakerneting some td files and relying on loose IR compatibility). Running through IREE requires the ability to import `iree.compiler` and `iree.runtime`, which can be obtained as follows: ``` python3 -m pip install iree-compiler-snapshot iree-runtime-snapshot -f https://github.com/google/iree/releases/tag/snapshot-20210406.200 PYTHONPATH="${PYTHONPATH}:${MY_IREE_BUILD}/bindings/python/" ``` This patch makes it painfully clear that we don't have any e2e testing harness to really plug into, and also don't have a usable Python API to our compiler stack (something usable in a jupyter notebook). That will be addressed in subsequent commits. We've been flying by the seat of our pants with this `examples` directory that isn't subject to any kind of testing or real usability concerns.
2021-04-09 04:05:16 +08:00
mlir::NPCOMP::IREEBackend::registerIREEBackendPasses();
}