This commit is contained in:
NiccoloN
@@ -0,0 +1,238 @@
|
||||
#include "mlir/Dialect/Func/IR/FuncOps.h"
|
||||
#include "mlir/IR/BuiltinAttributes.h"
|
||||
#include "mlir/IR/Diagnostics.h"
|
||||
#include "mlir/Support/LLVM.h"
|
||||
|
||||
#include "Common/IR/WeightUtils.hpp"
|
||||
#include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
|
||||
|
||||
using namespace mlir;
|
||||
|
||||
namespace onnx_mlir {
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr StringLiteral kPhaseMarker = "phase-check";
|
||||
|
||||
void checkWeightUseChains(func::FuncOp func, pim::CappedDiagnosticReporter& diagnostics) {
|
||||
func.walk([&](Operation* op) {
|
||||
if (!hasWeightAlways(op))
|
||||
return;
|
||||
|
||||
for (Value result : op->getResults()) {
|
||||
if (hasOnlySpatialMvmVmmWeightUses(result))
|
||||
continue;
|
||||
|
||||
diagnostics.report(op, [&](Operation* illegalOp) {
|
||||
illegalOp->emitOpError()
|
||||
<< kPhaseMarker
|
||||
<< " weight-marked values may only flow through static view/slice helper chains into Spatial VMM weights";
|
||||
});
|
||||
return;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
bool isRegionOrAncestorOf(Region& region, Region* candidate) {
|
||||
return candidate && (®ion == candidate || region.isAncestor(candidate));
|
||||
}
|
||||
|
||||
bool isValueDefinedInsideRegion(Value value, Region& region) {
|
||||
if (auto blockArg = dyn_cast<BlockArgument>(value))
|
||||
return isRegionOrAncestorOf(region, blockArg.getOwner()->getParent());
|
||||
if (Operation* definingOp = value.getDefiningOp())
|
||||
return isRegionOrAncestorOf(region, definingOp->getParentRegion());
|
||||
return false;
|
||||
}
|
||||
|
||||
bool isLegalExternalCapture(Value value, Region& region) {
|
||||
if (isValueDefinedInsideRegion(value, region))
|
||||
return true;
|
||||
|
||||
Operation* definingOp = value.getDefiningOp();
|
||||
return definingOp && definingOp->hasTrait<OpTrait::ConstantLike>();
|
||||
}
|
||||
|
||||
bool isRecordedDeferredCommunicationSource(Operation* op, Value value) {
|
||||
auto transfer = dyn_cast<spatial::SpatDeferredCommunicationOp>(op);
|
||||
return transfer && llvm::is_contained(transfer.getSources(), value);
|
||||
}
|
||||
|
||||
template <typename ComputeOpTy>
|
||||
void verifyComputeBodyCaptures(ComputeOpTy compute, StringRef kind, pim::CappedDiagnosticReporter& diagnostics) {
|
||||
Region& body = compute.getBody();
|
||||
body.walk([&](Operation* nestedOp) {
|
||||
for (OpOperand& operand : nestedOp->getOpOperands()) {
|
||||
Value value = operand.get();
|
||||
if (isLegalExternalCapture(value, body) || isRecordedDeferredCommunicationSource(nestedOp, value))
|
||||
continue;
|
||||
|
||||
Operation* definingOp = value.getDefiningOp();
|
||||
diagnostics.report(compute.getOperation(), [&](Operation* illegalOp) {
|
||||
InFlightDiagnostic diag =
|
||||
illegalOp->emitOpError() << kPhaseMarker << " " << kind << " body captures non-constant external operand #"
|
||||
<< operand.getOperandNumber() << " used by " << nestedOp->getName().getStringRef();
|
||||
diag << " (type " << value.getType() << ")";
|
||||
if (definingOp)
|
||||
diag.attachNote(definingOp->getLoc()) << "defining op is " << definingOp->getName().getStringRef();
|
||||
else if (auto blockArg = dyn_cast<BlockArgument>(value)) {
|
||||
if (Operation* owner = blockArg.getOwner()->getParentOp())
|
||||
diag.attachNote(owner->getLoc())
|
||||
<< "external block argument belongs to " << owner->getName().getStringRef();
|
||||
}
|
||||
});
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
bool isLegalHostBackedValue(Value value) {
|
||||
Operation* definingOp = value.getDefiningOp();
|
||||
if (!definingOp)
|
||||
return isa<BlockArgument>(value);
|
||||
return definingOp->getDialect()->getNamespace() != "spat";
|
||||
}
|
||||
|
||||
bool isScheduledPhase1Value(Value value) {
|
||||
Operation* definingOp = value.getDefiningOp();
|
||||
return isa_and_nonnull<spatial::SpatScheduledCompute, spatial::SpatScheduledComputeBatch>(definingOp);
|
||||
}
|
||||
|
||||
template <typename ComputeOpTy>
|
||||
void verifyScheduledInputs(ComputeOpTy compute,
|
||||
bool allowChannelReceiveInputs,
|
||||
StringRef kind,
|
||||
pim::CappedDiagnosticReporter& diagnostics) {
|
||||
for (auto [inputIndex, input] : llvm::enumerate(compute.getInputs())) {
|
||||
size_t currentInputIndex = inputIndex;
|
||||
Operation* definingOp = input.getDefiningOp();
|
||||
if (allowChannelReceiveInputs && isa_and_nonnull<spatial::SpatChannelReceiveOp>(definingOp))
|
||||
continue;
|
||||
if (isScheduledPhase1Value(input))
|
||||
continue;
|
||||
if (isLegalHostBackedValue(input))
|
||||
continue;
|
||||
|
||||
diagnostics.report(compute.getOperation(), [&](Operation* illegalOp) {
|
||||
InFlightDiagnostic diag = illegalOp->emitOpError()
|
||||
<< kPhaseMarker << " " << kind << " input #" << currentInputIndex
|
||||
<< (allowChannelReceiveInputs ? " must come from the host or explicit spat.channel_receive"
|
||||
: " must come from the host");
|
||||
if (definingOp)
|
||||
diag.attachNote(definingOp->getLoc()) << "illegal producer is " << definingOp->getName().getStringRef();
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
template <typename ComputeOpTy>
|
||||
void verifyNoNestedFragmentAssemblyBlueprints(ComputeOpTy compute,
|
||||
pim::CappedDiagnosticReporter& diagnostics) {
|
||||
compute.getBody().walk([&](spatial::SpatBlueprintOp blueprint) {
|
||||
if (!spatial::isFragmentAssembly(blueprint.getMode()))
|
||||
return;
|
||||
diagnostics.report(blueprint.getOperation(), [&](Operation* illegalOp) {
|
||||
illegalOp->emitOpError("fragment assembly blueprint must be host-level after merge materialization");
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter& diagnostics) {
|
||||
for (Operation& op : funcOp.getOps()) {
|
||||
if (isa<func::ReturnOp,
|
||||
spatial::SpatGraphCompute,
|
||||
spatial::SpatGraphComputeBatch,
|
||||
spatial::SpatConv2DPlanOp,
|
||||
spatial::SpatBiasAddPlanOp,
|
||||
spatial::SpatAddPlanOp,
|
||||
spatial::SpatConcatPlanOp,
|
||||
spatial::SpatFlattenPlanOp,
|
||||
spatial::SpatReluPlanOp,
|
||||
spatial::SpatSiluPlanOp,
|
||||
spatial::SpatResizeNearestPlanOp,
|
||||
spatial::SpatMaxPool2DPlanOp,
|
||||
spatial::SpatGlobalAveragePoolPlanOp,
|
||||
spatial::SpatBlueprintOp,
|
||||
spatial::SpatMaterializeLayoutOp>(&op)) {
|
||||
continue;
|
||||
}
|
||||
if (isa<spatial::SpatScheduledCompute, spatial::SpatScheduledComputeBatch>(&op)) {
|
||||
diagnostics.report(&op, [&](Operation* illegalOp) {
|
||||
illegalOp->emitOpError() << kPhaseMarker << " scheduled Spatial compute op is not allowed in logical graph phase";
|
||||
});
|
||||
continue;
|
||||
}
|
||||
if (isa<spatial::SpatChannelReceiveOp, spatial::SpatChannelSendOp>(&op)) {
|
||||
diagnostics.report(&op, [&](Operation* illegalOp) {
|
||||
illegalOp->emitOpError() << kPhaseMarker
|
||||
<< " explicit channel communication is not expected before merge materialization";
|
||||
});
|
||||
continue;
|
||||
}
|
||||
if (isCompileTimeOp(&op))
|
||||
continue;
|
||||
|
||||
diagnostics.report(&op, [&](Operation* illegalOp) {
|
||||
illegalOp->emitOpError()
|
||||
<< kPhaseMarker << " non-foldable top-level runtime op remains in logical Spatial graph; lower it inside spat.graph_compute";
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
void verifyScheduledTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter& diagnostics) {
|
||||
for (Operation& op : funcOp.getOps()) {
|
||||
if (isa<spatial::SpatChannelSendOp, spatial::SpatChannelReceiveOp>(&op)) {
|
||||
diagnostics.report(&op, [&](Operation* illegalOp) {
|
||||
illegalOp->emitOpError() << kPhaseMarker << " real channel communication is not allowed in scheduled phase 1";
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
LogicalResult verifyNoComputeBodyCaptures(func::FuncOp funcOp) {
|
||||
pim::CappedDiagnosticReporter diagnostics;
|
||||
for (auto compute : funcOp.getOps<spatial::SpatGraphCompute>())
|
||||
verifyComputeBodyCaptures(compute, "graph_compute", diagnostics);
|
||||
for (auto batch : funcOp.getOps<spatial::SpatGraphComputeBatch>())
|
||||
verifyComputeBodyCaptures(batch, "graph_compute_batch", diagnostics);
|
||||
for (auto compute : funcOp.getOps<spatial::SpatScheduledCompute>())
|
||||
verifyComputeBodyCaptures(compute, "scheduled_compute", diagnostics);
|
||||
for (auto batch : funcOp.getOps<spatial::SpatScheduledComputeBatch>())
|
||||
verifyComputeBodyCaptures(batch, "scheduled_compute_batch", diagnostics);
|
||||
diagnostics.emitSuppressedSummary(funcOp, "compute body capture verification failed");
|
||||
return success(!diagnostics.hasFailure());
|
||||
}
|
||||
|
||||
LogicalResult verifyONNXToSpatial(func::FuncOp funcOp) { return verifyLogicalSpatialGraphInvariants(funcOp); }
|
||||
|
||||
LogicalResult verifyLogicalSpatialGraphInvariants(func::FuncOp funcOp) {
|
||||
pim::CappedDiagnosticReporter diagnostics;
|
||||
verifyLogicalTopLevelOps(funcOp, diagnostics);
|
||||
checkWeightUseChains(funcOp, diagnostics);
|
||||
if (failed(verifyNoComputeBodyCaptures(funcOp)))
|
||||
return failure();
|
||||
diagnostics.emitSuppressedSummary(funcOp, "logical Spatial graph verification failed");
|
||||
return success(!diagnostics.hasFailure());
|
||||
}
|
||||
|
||||
LogicalResult verifyScheduledSpatialInvariants(func::FuncOp funcOp) {
|
||||
pim::CappedDiagnosticReporter diagnostics;
|
||||
verifyScheduledTopLevelOps(funcOp, diagnostics);
|
||||
for (auto compute : funcOp.getOps<spatial::SpatScheduledCompute>()) {
|
||||
verifyScheduledInputs(compute, /*allowChannelReceiveInputs=*/true, "spat.scheduled_compute", diagnostics);
|
||||
verifyNoNestedFragmentAssemblyBlueprints(compute, diagnostics);
|
||||
}
|
||||
for (auto batch : funcOp.getOps<spatial::SpatScheduledComputeBatch>()) {
|
||||
verifyScheduledInputs(batch, /*allowChannelReceiveInputs=*/false, "spat.scheduled_compute_batch", diagnostics);
|
||||
verifyNoNestedFragmentAssemblyBlueprints(batch, diagnostics);
|
||||
}
|
||||
if (failed(verifyNoComputeBodyCaptures(funcOp)))
|
||||
return failure();
|
||||
diagnostics.emitSuppressedSummary(funcOp, "scheduled Spatial verification failed");
|
||||
return success(!diagnostics.hasFailure());
|
||||
}
|
||||
|
||||
} // namespace onnx_mlir
|
||||
@@ -0,0 +1,13 @@
|
||||
#pragma once
|
||||
|
||||
#include "mlir/Dialect/Func/IR/FuncOps.h"
|
||||
#include "mlir/Support/LogicalResult.h"
|
||||
|
||||
namespace onnx_mlir {
|
||||
|
||||
mlir::LogicalResult verifyONNXToSpatial(mlir::func::FuncOp funcOp);
|
||||
mlir::LogicalResult verifyNoComputeBodyCaptures(mlir::func::FuncOp funcOp);
|
||||
mlir::LogicalResult verifyLogicalSpatialGraphInvariants(mlir::func::FuncOp funcOp);
|
||||
mlir::LogicalResult verifyScheduledSpatialInvariants(mlir::func::FuncOp funcOp);
|
||||
|
||||
} // namespace onnx_mlir
|
||||
@@ -0,0 +1,152 @@
|
||||
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Transforms/PlanLowering.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
|
||||
|
||||
using namespace mlir;
|
||||
|
||||
namespace onnx_mlir::spatial {
|
||||
|
||||
static LayoutAlternative denseAlternative(Operation *op) {
|
||||
LayoutAlternative alternative;
|
||||
alternative.operandLayouts.assign(op->getNumOperands(), PhysicalLayout::DenseNCHW);
|
||||
alternative.resultLayout = PhysicalLayout::DenseNCHW;
|
||||
return alternative;
|
||||
}
|
||||
|
||||
static LayoutAlternative rowStripAlternative(Operation *op,
|
||||
ArrayRef<PhysicalLayout> operandLayouts) {
|
||||
LayoutAlternative alternative;
|
||||
alternative.operandLayouts.assign(operandLayouts.begin(), operandLayouts.end());
|
||||
alternative.resultLayout = PhysicalLayout::NHWCRowStrip;
|
||||
alternative.intrinsicCost = -2;
|
||||
return alternative;
|
||||
}
|
||||
|
||||
static bool hasRowStripInput(ArrayRef<PhysicalLayout> operandLayouts, unsigned index) {
|
||||
return index < operandLayouts.size()
|
||||
&& operandLayouts[index] == PhysicalLayout::NHWCRowStrip;
|
||||
}
|
||||
|
||||
SmallVector<LayoutAlternative> SpatConv2DPlanOp::getLayoutAlternatives(
|
||||
const SpatialTargetResources& target, ArrayRef<PhysicalLayout> operandLayouts) {
|
||||
SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
|
||||
if (hasRowStripInput(operandLayouts, 0)) {
|
||||
if (succeeded(canConsumeAndProduceRowStrip(*this, target)))
|
||||
alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts));
|
||||
}
|
||||
else if (succeeded(canLowerConvPlanToRowStrip(*this, target))) {
|
||||
LayoutAlternative alternative = denseAlternative(getOperation());
|
||||
alternative.resultLayout = PhysicalLayout::NHWCRowStrip;
|
||||
alternative.intrinsicCost = -2;
|
||||
alternatives.push_back(std::move(alternative));
|
||||
}
|
||||
return alternatives;
|
||||
}
|
||||
|
||||
SmallVector<LayoutAlternative> SpatFlattenPlanOp::getLayoutAlternatives(
|
||||
const SpatialTargetResources& target, ArrayRef<PhysicalLayout> operandLayouts) {
|
||||
SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
|
||||
if (!operandLayouts.empty()
|
||||
&& operandLayouts[0] == PhysicalLayout::Fragmented) {
|
||||
LayoutAlternative alternative = denseAlternative(getOperation());
|
||||
alternative.operandLayouts[0] = PhysicalLayout::Fragmented;
|
||||
alternatives.push_back(std::move(alternative));
|
||||
}
|
||||
if (hasRowStripInput(operandLayouts, 0)
|
||||
&& succeeded(canLowerFlattenFromRowStrip(*this, target))) {
|
||||
LayoutAlternative alternative = rowStripAlternative(getOperation(), operandLayouts);
|
||||
alternative.resultLayout = PhysicalLayout::DenseNCHW;
|
||||
alternatives.push_back(std::move(alternative));
|
||||
}
|
||||
return alternatives;
|
||||
}
|
||||
|
||||
SmallVector<LayoutAlternative> SpatReluPlanOp::getLayoutAlternatives(
|
||||
const SpatialTargetResources&, ArrayRef<PhysicalLayout> operandLayouts) {
|
||||
SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
|
||||
if (hasRowStripInput(operandLayouts, 0))
|
||||
alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts));
|
||||
return alternatives;
|
||||
}
|
||||
|
||||
SmallVector<LayoutAlternative> SpatSiluPlanOp::getLayoutAlternatives(
|
||||
const SpatialTargetResources&, ArrayRef<PhysicalLayout> operandLayouts) {
|
||||
SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
|
||||
if (hasRowStripInput(operandLayouts, 0)) {
|
||||
LayoutAlternative alternative = rowStripAlternative(getOperation(), operandLayouts);
|
||||
alternative.intrinsicCost = -3;
|
||||
alternatives.push_back(std::move(alternative));
|
||||
}
|
||||
return alternatives;
|
||||
}
|
||||
|
||||
SmallVector<LayoutAlternative> SpatResizeNearestPlanOp::getLayoutAlternatives(
|
||||
const SpatialTargetResources& target, ArrayRef<PhysicalLayout> operandLayouts) {
|
||||
SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
|
||||
if (hasRowStripInput(operandLayouts, 0)
|
||||
&& succeeded(canLowerResizeNearestPlanToRowStrip(*this, target)))
|
||||
alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts));
|
||||
return alternatives;
|
||||
}
|
||||
|
||||
SmallVector<LayoutAlternative> SpatMaxPool2DPlanOp::getLayoutAlternatives(
|
||||
const SpatialTargetResources& target, ArrayRef<PhysicalLayout> operandLayouts) {
|
||||
SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
|
||||
if (succeeded(canLowerMaxPoolPlanToRowStrip(*this, target))) {
|
||||
LayoutAlternative alternative = denseAlternative(getOperation());
|
||||
if (hasRowStripInput(operandLayouts, 0))
|
||||
alternative = rowStripAlternative(getOperation(), operandLayouts);
|
||||
alternative.resultLayout = PhysicalLayout::NHWCRowStrip;
|
||||
alternative.intrinsicCost = -2;
|
||||
alternatives.push_back(std::move(alternative));
|
||||
}
|
||||
return alternatives;
|
||||
}
|
||||
|
||||
SmallVector<LayoutAlternative> SpatGlobalAveragePoolPlanOp::getLayoutAlternatives(
|
||||
const SpatialTargetResources& target, ArrayRef<PhysicalLayout> operandLayouts) {
|
||||
SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
|
||||
if (succeeded(canLowerGlobalAveragePoolPlanToRowStrip(*this, target))) {
|
||||
LayoutAlternative alternative = denseAlternative(getOperation());
|
||||
if (hasRowStripInput(operandLayouts, 0))
|
||||
alternative = rowStripAlternative(getOperation(), operandLayouts);
|
||||
alternative.resultLayout = PhysicalLayout::NHWCRowStrip;
|
||||
alternative.intrinsicCost = -2;
|
||||
alternatives.push_back(std::move(alternative));
|
||||
}
|
||||
return alternatives;
|
||||
}
|
||||
|
||||
SmallVector<LayoutAlternative> SpatBiasAddPlanOp::getLayoutAlternatives(
|
||||
const SpatialTargetResources&, ArrayRef<PhysicalLayout> operandLayouts) {
|
||||
SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
|
||||
auto resultType = dyn_cast<RankedTensorType>(getOutput().getType());
|
||||
if (resultType && hasRowStripInput(operandLayouts, 0)
|
||||
&& isSupportedBiasAddValue(getBias(), resultType))
|
||||
alternatives.push_back(rowStripAlternative(getOperation(),
|
||||
{PhysicalLayout::NHWCRowStrip,
|
||||
PhysicalLayout::DenseNCHW}));
|
||||
return alternatives;
|
||||
}
|
||||
|
||||
SmallVector<LayoutAlternative> SpatAddPlanOp::getLayoutAlternatives(
|
||||
const SpatialTargetResources&, ArrayRef<PhysicalLayout> operandLayouts) {
|
||||
SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
|
||||
if (operandLayouts.size() >= 2 && hasRowStripInput(operandLayouts, 0)
|
||||
&& hasRowStripInput(operandLayouts, 1))
|
||||
alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts));
|
||||
return alternatives;
|
||||
}
|
||||
|
||||
SmallVector<LayoutAlternative> SpatConcatPlanOp::getLayoutAlternatives(
|
||||
const SpatialTargetResources&, ArrayRef<PhysicalLayout> operandLayouts) {
|
||||
SmallVector<LayoutAlternative> alternatives {denseAlternative(getOperation())};
|
||||
if (!operandLayouts.empty() && llvm::all_of(operandLayouts, [](PhysicalLayout layout) {
|
||||
return layout == PhysicalLayout::NHWCRowStrip;
|
||||
}))
|
||||
alternatives.push_back(rowStripAlternative(getOperation(), operandLayouts));
|
||||
return alternatives;
|
||||
}
|
||||
|
||||
} // namespace onnx_mlir::spatial
|
||||
Reference in New Issue
Block a user