torch-mlir/include/npcomp/RefBackend/RefBackend.h

71 lines
2.5 KiB
C
Raw Normal View History

//===------------------------------------------------------------*- 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_REFBACKEND_REFBACKEND_H
#define NPCOMP_REFBACKEND_REFBACKEND_H
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassManager.h"
namespace mlir {
namespace NPCOMP {
/// Registers all RefBackend passes.
void registerRefBackendPasses();
// Look in createRefBackendLoweringPipeline for more information about how these
// passes fit together.
//
// Pass summaries are in Passes.td.
Totally rework RefE2E tensor to memref flow. (#42) This now gets the overall "RefE2E" compilation stack to a point that I'm fairly happy with. We simplify it by mostly embracing the "descriptor" view of the world. The overall flow is best understood by reading through the createE2ELoweringPipeline function in lib/E2E/E2E.cpp That function creates a pass pipeline that lowers from "TCF" (which is ~numpy level of abstraction) down to LLVM IR. A brief high-level summary of what happens there: 1. TCF to TCP conversion. This involves reifying error handling in the form of shape constraints. See test/Conversion/TCFToTCP/basic.mlir 2. Lowering shape constraints. This converts shape constraints into eager error-handling code. See test/E2E/lower-shape-constraints.mlir This pass will soon go upstream. Because this lowers to std.assert, some later passes like LowerToNpcomprtABI and LowerToLLVM are updated to properly plumb this through e2e. See test/npcomp-run-mlir/invalid-broadcast.mlir for an execution test that properly aborts in case of an error. 3. Lowering tensors to memrefs. This is done via a series of passes rather than an single mega conversion. Unlike the previous code that mixed in the npcomprt ABI stuff here, it's now a very clean "pure memref" conversion. See test/E2E/lower-*-to-memref.mlir and lib/E2E/TensorToMemref/ Most of the changes are concentrated here. 4. As part of the above, we use the upstream ConvertShapeToStandard for lowering shapes. 5. We lower linalg to loops and lower loops to CFG using upstream passes. 6. Rewrite the "ABI" boundaries of the program to npcomprt data structures (LowerToNpcomprtABI). This mainly affects ABI boundaries and how global tensor constants are represented. One of the major improvements in this commit is that now it's a very clean rewrite that just replaces memrefs on ABI boundaries with !npcomprt.tensor (before there was a get_extent function that is not needed). See test/E2E/lower-to-npcomprt-abi.mlir 7. Lower to LLVM with upstream mlir patterns + some patterns for the npcomprt lowerings. One aspect here that is still a remnant of a non-descriptor-based tensor to memref flow is the BypassShapes + LowerShapedResultsToMemref. BypassShapes wraps the "tensor compute" ops in a tcp.shaped_results (basically a "tie_shape" kind of op), and then LowerShapedResultsToMemref uses those annotations to allocate output buffers while lowering the "tensor compute ops". Note that there are very few "tensor compute" ops currently supported (tcp.add + tcp.broadcast_to), so we just hardcode them in both passes. Realistically, I expect this to go away as we fully embrace the descriptor-based approach for simplicity, so don't look too deep into it.
2020-09-17 08:31:40 +08:00
std::unique_ptr<OperationPass<FuncOp>> createLowerStructuralToMemrefPass();
std::unique_ptr<OperationPass<ModuleOp>> createLowerToRefbackrtABIPass();
std::unique_ptr<OperationPass<FuncOp>> createLowerAllocMemRefOpsPass();
std::unique_ptr<OperationPass<ModuleOp>> createLowerToLLVMPass();
std::unique_ptr<Pass> createRestrictedCanonicalizerPass();
struct RefBackendLoweringPipelineOptions
: public PassPipelineOptions<RefBackendLoweringPipelineOptions> {
// If this option is true, then perform optimizations.
// If this option is false, only do the bare minimum for correctness.
Option<bool> optimize{*this, "optimize", llvm::cl::desc("Do optimizations."),
llvm::cl::init(false)};
};
// The main pipeline that encapsulates the full RefBackend lowering.
void createRefBackendLoweringPipeline(
OpPassManager &pm, const RefBackendLoweringPipelineOptions &options);
// Helper pipeline that runs TCF->TCP lowering.
//
// For now, just piggy-back on the same set of options since this is such a
// simple set of passes.
//
// TODO: Move this out of RefBackend once the TCF->TCP conversions
// become more substantial.
void createRefBackendTCFToTCPPipeline(
OpPassManager &pm, const RefBackendLoweringPipelineOptions &options);
// Helper pipeline that runs TCF->TCP lowering before invoking
// RefBackendLoweringPipeline.
// For now, just piggy-back on the same set of options since this is such a
// thin wrapper.
// Longer-term, the reference backend should fit into some sort of
// "target interface" and this helper won't be needed.
void createTCFRefBackendLoweringPipeline(
OpPassManager &pm, const RefBackendLoweringPipelineOptions &options);
} // namespace NPCOMP
} // namespace mlir
#endif // NPCOMP_REFBACKEND_REFBACKEND_H