2020-09-17 08:31:40 +08:00
|
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
//
|
|
|
|
// Part of the LLVM Project, 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 "PassDetail.h"
|
|
|
|
#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
|
|
|
|
#include "mlir/Dialect/Shape/IR/Shape.h"
|
2020-09-18 09:56:01 +08:00
|
|
|
#include "mlir/Dialect/StandardOps/IR/Ops.h"
|
2020-10-08 08:30:10 +08:00
|
|
|
#include "npcomp/Dialect/Refback/IR/RefbackDialect.h"
|
|
|
|
#include "npcomp/Dialect/Refback/IR/RefbackOps.h"
|
2020-09-17 08:31:40 +08:00
|
|
|
#include "npcomp/Dialect/TCP/IR/TCPOps.h"
|
2020-10-07 07:14:37 +08:00
|
|
|
#include "npcomp/RefBackend/RefBackend.h"
|
2020-09-17 08:31:40 +08:00
|
|
|
|
|
|
|
using namespace mlir;
|
|
|
|
using namespace mlir::NPCOMP;
|
|
|
|
|
|
|
|
// TODO: Don't just open-code all shape transfer functions here.
|
|
|
|
static SmallVector<Value, 6> bypassResultShapes(Operation &op) {
|
|
|
|
OpBuilder builder(&op);
|
|
|
|
|
|
|
|
if (auto broadcastTo = dyn_cast<tcp::BroadcastToOp>(op)) {
|
|
|
|
return {broadcastTo.shape()};
|
|
|
|
}
|
|
|
|
|
2020-09-25 08:14:21 +08:00
|
|
|
// Elementwise ops.
|
|
|
|
if (isa<tcp::AddOp, tcp::MaxOp, tcp::ExpOp, tcp::TanhOp>(op)) {
|
2020-09-19 03:53:55 +08:00
|
|
|
return {builder.create<shape::ShapeOfOp>(op.getLoc(), op.getOperand(0))};
|
|
|
|
}
|
|
|
|
|
2020-09-18 09:56:01 +08:00
|
|
|
if (auto matmul = dyn_cast<tcp::MatmulOp>(op)) {
|
|
|
|
auto lhsRows = builder.create<DimOp>(op.getLoc(), matmul.lhs(), 0);
|
|
|
|
auto rhsCols = builder.create<DimOp>(op.getLoc(), matmul.rhs(), 1);
|
|
|
|
auto shape = builder.create<TensorFromElementsOp>(
|
|
|
|
op.getLoc(), ValueRange({lhsRows, rhsCols}));
|
|
|
|
return {shape};
|
|
|
|
}
|
|
|
|
|
2020-09-17 08:31:40 +08:00
|
|
|
// No shape transfer function.
|
|
|
|
return {};
|
|
|
|
}
|
|
|
|
|
|
|
|
namespace {
|
|
|
|
// TODO: There is a coupling between this pass and LowerShapedResults.
|
[RefBackend] Split out RefBackend (refback) dialect from TCP.
This is the first in a patch series that is refactoring the
constellation of things variously called or associated with "E2E",
"RefE2E", "npcomprt", and "TCP" into a more cleanly layered result.
Concretely, this first patch fixes the fact that TCP was basically
acting like a dumping ground needed by the reference backend. This
splits it out, which is fairly mechanical, but touches a lot of lines of
code (basically replacing `tcp` with `refback` and `TCP` with
`RefBackend).
Now, the RefBackend dialect is that dumping ground, which
is slighly better, as it starts allowing TCP to become a nice clean
middle layer that is not related per se to the reference backend.
The previous name RefE2E or "reference e2e flow" was super confusing.
Now that we are seeing more clearly where the "backend" distinction
lies, the [RefBackend] commit tag is born :)
2020-10-07 06:44:18 +08:00
|
|
|
// Any op that is wrapped in refback.shaped_results here needs to be known how
|
|
|
|
// to be lowered by LowerShapedResults.
|
2020-09-17 08:31:40 +08:00
|
|
|
class BypassShapes : public BypassShapesBase<BypassShapes> {
|
2020-09-22 05:48:44 +08:00
|
|
|
void getDependentDialects(DialectRegistry ®istry) const override {
|
2020-10-08 08:30:10 +08:00
|
|
|
registry.insert<shape::ShapeDialect, refback::RefbackDialect>();
|
2020-09-22 05:48:44 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
void runOnOperation() override {
|
2020-09-17 08:31:40 +08:00
|
|
|
auto func = getOperation();
|
|
|
|
func.walk([&](Operation *opPtr) {
|
|
|
|
Operation &op = *opPtr;
|
|
|
|
SmallVector<Value, 6> resultShapes = bypassResultShapes(op);
|
|
|
|
if (resultShapes.empty())
|
|
|
|
return;
|
[RefBackend] Split out RefBackend (refback) dialect from TCP.
This is the first in a patch series that is refactoring the
constellation of things variously called or associated with "E2E",
"RefE2E", "npcomprt", and "TCP" into a more cleanly layered result.
Concretely, this first patch fixes the fact that TCP was basically
acting like a dumping ground needed by the reference backend. This
splits it out, which is fairly mechanical, but touches a lot of lines of
code (basically replacing `tcp` with `refback` and `TCP` with
`RefBackend).
Now, the RefBackend dialect is that dumping ground, which
is slighly better, as it starts allowing TCP to become a nice clean
middle layer that is not related per se to the reference backend.
The previous name RefE2E or "reference e2e flow" was super confusing.
Now that we are seeing more clearly where the "backend" distinction
lies, the [RefBackend] commit tag is born :)
2020-10-07 06:44:18 +08:00
|
|
|
// We have result shapes, so wrap this op in a refback.shaped_results op.
|
2020-09-17 08:31:40 +08:00
|
|
|
OpBuilder builder(&op);
|
[RefBackend] Split out RefBackend (refback) dialect from TCP.
This is the first in a patch series that is refactoring the
constellation of things variously called or associated with "E2E",
"RefE2E", "npcomprt", and "TCP" into a more cleanly layered result.
Concretely, this first patch fixes the fact that TCP was basically
acting like a dumping ground needed by the reference backend. This
splits it out, which is fairly mechanical, but touches a lot of lines of
code (basically replacing `tcp` with `refback` and `TCP` with
`RefBackend).
Now, the RefBackend dialect is that dumping ground, which
is slighly better, as it starts allowing TCP to become a nice clean
middle layer that is not related per se to the reference backend.
The previous name RefE2E or "reference e2e flow" was super confusing.
Now that we are seeing more clearly where the "backend" distinction
lies, the [RefBackend] commit tag is born :)
2020-10-07 06:44:18 +08:00
|
|
|
auto shapedResults = builder.create<refback::ShapedResultsOp>(
|
2020-09-17 08:31:40 +08:00
|
|
|
op.getLoc(), op.getResultTypes(), resultShapes);
|
|
|
|
op.replaceAllUsesWith(shapedResults);
|
|
|
|
|
|
|
|
// Move the op into the body and yield the results.
|
|
|
|
Block *body = builder.createBlock(&shapedResults.body());
|
|
|
|
op.moveBefore(body, body->end());
|
[RefBackend] Split out RefBackend (refback) dialect from TCP.
This is the first in a patch series that is refactoring the
constellation of things variously called or associated with "E2E",
"RefE2E", "npcomprt", and "TCP" into a more cleanly layered result.
Concretely, this first patch fixes the fact that TCP was basically
acting like a dumping ground needed by the reference backend. This
splits it out, which is fairly mechanical, but touches a lot of lines of
code (basically replacing `tcp` with `refback` and `TCP` with
`RefBackend).
Now, the RefBackend dialect is that dumping ground, which
is slighly better, as it starts allowing TCP to become a nice clean
middle layer that is not related per se to the reference backend.
The previous name RefE2E or "reference e2e flow" was super confusing.
Now that we are seeing more clearly where the "backend" distinction
lies, the [RefBackend] commit tag is born :)
2020-10-07 06:44:18 +08:00
|
|
|
builder.create<refback::YieldOp>(op.getLoc(), op.getResults());
|
2020-09-17 08:31:40 +08:00
|
|
|
});
|
|
|
|
}
|
|
|
|
};
|
|
|
|
} // namespace
|
|
|
|
|
|
|
|
std::unique_ptr<OperationPass<FuncOp>> mlir::NPCOMP::createBypassShapesPass() {
|
|
|
|
return std::make_unique<BypassShapes>();
|
|
|
|
}
|