DeadLock
This commit is contained in:
ilgeco
@@ -0,0 +1,200 @@
|
||||
#include "mlir/Dialect/Func/IR/FuncOps.h"
|
||||
#include "mlir/IR/PatternMatch.h"
|
||||
#include "mlir/Pass/Pass.h"
|
||||
|
||||
#include "llvm/ADT/DenseMap.h"
|
||||
|
||||
#include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
|
||||
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
|
||||
#include "src/Accelerators/PIM/Pass/PIMPasses.h"
|
||||
#include "src/Accelerators/PIM/Pass/PIMPasses.h"
|
||||
|
||||
using namespace mlir;
|
||||
|
||||
namespace onnx_mlir {
|
||||
namespace {
|
||||
|
||||
static constexpr StringLiteral kLogicalLayout = "nchw";
|
||||
static constexpr StringLiteral kDenseLayout = "dense_nchw";
|
||||
static constexpr StringLiteral kRowStripLayout = "nchw_row_strip";
|
||||
static constexpr StringLiteral kRowStripIndexMap = "packed_hwc_rows_to_nchw";
|
||||
|
||||
enum class SelectedLayout {
|
||||
DenseNchw,
|
||||
NchwRowStrip,
|
||||
};
|
||||
|
||||
static SelectedLayout getSelectedLayout(llvm::DenseMap<Value, SelectedLayout>& layouts, Value value) {
|
||||
auto it = layouts.find(value);
|
||||
return it == layouts.end() ? SelectedLayout::DenseNchw : it->second;
|
||||
}
|
||||
|
||||
static bool usesSelectedRowStrip(Operation* user, llvm::DenseMap<Value, SelectedLayout>& layouts) {
|
||||
if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(user))
|
||||
return getSelectedLayout(layouts, reluPlan.getResult()) == SelectedLayout::NchwRowStrip;
|
||||
if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(user))
|
||||
return getSelectedLayout(layouts, convPlan.getResult()) == SelectedLayout::NchwRowStrip;
|
||||
return false;
|
||||
}
|
||||
|
||||
static bool allUsersCanHandleRowStrip(Value value, llvm::DenseMap<Value, SelectedLayout>& layouts) {
|
||||
for (Operation* user : value.getUsers()) {
|
||||
if (usesSelectedRowStrip(user, layouts))
|
||||
continue;
|
||||
// Dense-only users must be materialized explicitly.
|
||||
continue;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
static std::pair<SmallVector<int64_t>, SmallVector<int64_t>> buildRowStripMetadata(RankedTensorType type) {
|
||||
SmallVector<int64_t> offsets;
|
||||
SmallVector<int64_t> sizes;
|
||||
const int64_t channels = type.getDimSize(1);
|
||||
const int64_t height = type.getDimSize(2);
|
||||
const int64_t width = type.getDimSize(3);
|
||||
offsets.reserve(height * 4);
|
||||
sizes.reserve(height * 4);
|
||||
for (int64_t row = 0; row < height; ++row) {
|
||||
offsets.append({0, 0, row, 0});
|
||||
sizes.append({1, channels, 1, width});
|
||||
}
|
||||
return {offsets, sizes};
|
||||
}
|
||||
|
||||
static bool canSelectConvRowStrip(spatial::SpatConv2DPlanOp convPlan,
|
||||
llvm::DenseMap<Value, SelectedLayout>& layouts) {
|
||||
SelectedLayout inputLayout = getSelectedLayout(layouts, convPlan.getInput());
|
||||
if (inputLayout == SelectedLayout::NchwRowStrip)
|
||||
return succeeded(canConsumeAndProduceRowStrip(convPlan));
|
||||
return succeeded(canLowerConvPlanToRowStrip(convPlan));
|
||||
}
|
||||
|
||||
static SelectedLayout chooseConvLayout(spatial::SpatConv2DPlanOp convPlan,
|
||||
llvm::DenseMap<Value, SelectedLayout>& layouts) {
|
||||
if (!canSelectConvRowStrip(convPlan, layouts))
|
||||
return SelectedLayout::DenseNchw;
|
||||
if (!allUsersCanHandleRowStrip(convPlan.getResult(), layouts))
|
||||
return SelectedLayout::DenseNchw;
|
||||
return SelectedLayout::NchwRowStrip;
|
||||
}
|
||||
|
||||
static SelectedLayout chooseReluLayout(spatial::SpatReluPlanOp reluPlan,
|
||||
llvm::DenseMap<Value, SelectedLayout>& layouts) {
|
||||
if (getSelectedLayout(layouts, reluPlan.getInput()) != SelectedLayout::NchwRowStrip)
|
||||
return SelectedLayout::DenseNchw;
|
||||
if (!allUsersCanHandleRowStrip(reluPlan.getResult(), layouts))
|
||||
return SelectedLayout::DenseNchw;
|
||||
return SelectedLayout::NchwRowStrip;
|
||||
}
|
||||
|
||||
static spatial::SpatReconciliatorOp insertRowStripReconciliator(IRRewriter& rewriter, Value value) {
|
||||
auto outputType = cast<RankedTensorType>(value.getType());
|
||||
auto [offsets, sizes] = buildRowStripMetadata(outputType);
|
||||
return spatial::SpatReconciliatorOp::create(rewriter,
|
||||
value.getLoc(),
|
||||
outputType,
|
||||
value,
|
||||
rewriter.getStringAttr(kLogicalLayout),
|
||||
rewriter.getStringAttr(kRowStripLayout),
|
||||
rewriter.getDenseI64ArrayAttr(offsets),
|
||||
rewriter.getDenseI64ArrayAttr(sizes),
|
||||
rewriter.getStringAttr(kRowStripIndexMap));
|
||||
}
|
||||
|
||||
static void materializeDenseUses(IRRewriter& rewriter,
|
||||
Value layoutValue,
|
||||
llvm::DenseMap<Value, SelectedLayout>& layouts) {
|
||||
SmallVector<OpOperand*> denseUses;
|
||||
for (OpOperand& use : layoutValue.getUses()) {
|
||||
if (usesSelectedRowStrip(use.getOwner(), layouts))
|
||||
continue;
|
||||
denseUses.push_back(&use);
|
||||
}
|
||||
|
||||
for (OpOperand* use : denseUses) {
|
||||
Operation* owner = use->getOwner();
|
||||
rewriter.setInsertionPoint(owner);
|
||||
auto materialized = spatial::SpatMaterializeLayoutOp::create(rewriter,
|
||||
owner->getLoc(),
|
||||
use->get().getType(),
|
||||
use->get(),
|
||||
rewriter.getStringAttr(kLogicalLayout),
|
||||
rewriter.getStringAttr(kRowStripLayout),
|
||||
rewriter.getStringAttr(kDenseLayout));
|
||||
use->set(materialized.getResult());
|
||||
}
|
||||
}
|
||||
|
||||
struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass, OperationPass<ModuleOp>> {
|
||||
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SpatialLayoutPlanningPass)
|
||||
|
||||
StringRef getArgument() const override { return "spatial-layout-planning"; }
|
||||
StringRef getDescription() const override { return "Select conservative Spatial layouts and insert reconciliation barriers."; }
|
||||
|
||||
void runOnOperation() override {
|
||||
auto entryFunc = getPimEntryFunc(getOperation());
|
||||
if (failed(entryFunc)) {
|
||||
getOperation().emitError("failed to locate the PIM entry function during Spatial layout planning");
|
||||
signalPassFailure();
|
||||
return;
|
||||
}
|
||||
|
||||
func::FuncOp funcOp = *entryFunc;
|
||||
IRRewriter rewriter(&getContext());
|
||||
llvm::DenseMap<Value, SelectedLayout> layouts;
|
||||
|
||||
bool changed = true;
|
||||
while (changed) {
|
||||
changed = false;
|
||||
for (Operation& op : llvm::make_early_inc_range(funcOp.getBody().front())) {
|
||||
if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(&op)) {
|
||||
SelectedLayout selected = chooseConvLayout(convPlan, layouts);
|
||||
if (layouts[convPlan.getResult()] != selected) {
|
||||
layouts[convPlan.getResult()] = selected;
|
||||
changed = true;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(&op)) {
|
||||
SelectedLayout selected = chooseReluLayout(reluPlan, layouts);
|
||||
if (layouts[reluPlan.getResult()] != selected) {
|
||||
layouts[reluPlan.getResult()] = selected;
|
||||
changed = true;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (Operation& op : llvm::make_early_inc_range(funcOp.getBody().front())) {
|
||||
Value producedValue;
|
||||
if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(&op))
|
||||
producedValue = convPlan.getResult();
|
||||
else if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(&op))
|
||||
producedValue = reluPlan.getResult();
|
||||
else
|
||||
continue;
|
||||
|
||||
if (getSelectedLayout(layouts, producedValue) != SelectedLayout::NchwRowStrip)
|
||||
continue;
|
||||
|
||||
rewriter.setInsertionPointAfter(&op);
|
||||
auto reconciliator = insertRowStripReconciliator(rewriter, producedValue);
|
||||
rewriter.replaceAllUsesExcept(producedValue, reconciliator.getResult(), reconciliator);
|
||||
materializeDenseUses(rewriter, reconciliator.getResult(), layouts);
|
||||
}
|
||||
if (failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) {
|
||||
getOperation().emitError("RAPTOR_PHASE_CHECK logical Spatial graph verification failed after SpatialLayoutPlanning");
|
||||
signalPassFailure();
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
std::unique_ptr<Pass> createSpatialLayoutPlanningPass() { return std::make_unique<SpatialLayoutPlanningPass>(); }
|
||||
|
||||
} // namespace onnx_mlir
|
||||
Reference in New Issue
Block a user