This commit is contained in:
NiccoloN
@@ -1,39 +0,0 @@
|
||||
#include "ContractionMaterialization.hpp"
|
||||
|
||||
#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
|
||||
#include "MatrixProductLowering.hpp"
|
||||
|
||||
namespace onnx_mlir {
|
||||
|
||||
mlir::Value materializePaddedContractionInput(
|
||||
mlir::Value input,
|
||||
mlir::RankedTensorType paddedType,
|
||||
mlir::PatternRewriter& rewriter,
|
||||
mlir::Location loc) {
|
||||
return createPaddedInputCompute(input, paddedType, rewriter, loc);
|
||||
}
|
||||
|
||||
mlir::FailureOr<mlir::Value> materializeTransposedContractionConstant(
|
||||
mlir::Value input,
|
||||
mlir::RankedTensorType resultType,
|
||||
llvm::ArrayRef<int64_t> permutation,
|
||||
mlir::PatternRewriter& rewriter,
|
||||
mlir::Location loc) {
|
||||
auto denseAttr = getHostConstDenseElementsAttr(input);
|
||||
auto inputType = denseAttr ? mlir::dyn_cast<mlir::RankedTensorType>(denseAttr.getType()) : nullptr;
|
||||
if (!inputType || !inputType.hasStaticShape() || !resultType || !resultType.hasStaticShape()
|
||||
|| inputType.getRank() != resultType.getRank())
|
||||
return mlir::failure();
|
||||
|
||||
auto transposedAttr = transposeDenseElementsAttr(denseAttr, permutation);
|
||||
if (mlir::failed(transposedAttr) || transposedAttr->getType() != resultType)
|
||||
return mlir::failure();
|
||||
|
||||
return getOrCreateConstant(rewriter,
|
||||
rewriter.getInsertionBlock()->getParentOp(),
|
||||
*transposedAttr,
|
||||
resultType);
|
||||
}
|
||||
|
||||
} // namespace onnx_mlir
|
||||
@@ -1,23 +0,0 @@
|
||||
#pragma once
|
||||
|
||||
#include "llvm/ADT/ArrayRef.h"
|
||||
|
||||
#include "mlir/IR/BuiltinTypes.h"
|
||||
#include "mlir/IR/PatternMatch.h"
|
||||
|
||||
namespace onnx_mlir {
|
||||
|
||||
mlir::Value materializePaddedContractionInput(
|
||||
mlir::Value input,
|
||||
mlir::RankedTensorType paddedType,
|
||||
mlir::PatternRewriter& rewriter,
|
||||
mlir::Location loc);
|
||||
|
||||
mlir::FailureOr<mlir::Value> materializeTransposedContractionConstant(
|
||||
mlir::Value input,
|
||||
mlir::RankedTensorType resultType,
|
||||
llvm::ArrayRef<int64_t> permutation,
|
||||
mlir::PatternRewriter& rewriter,
|
||||
mlir::Location loc);
|
||||
|
||||
} // namespace onnx_mlir
|
||||
@@ -1,902 +0,0 @@
|
||||
#include "mlir/Dialect/Affine/IR/AffineOps.h"
|
||||
#include "mlir/Dialect/Arith/IR/Arith.h"
|
||||
#include "mlir/Dialect/Func/IR/FuncOps.h"
|
||||
#include "mlir/Dialect/Linalg/IR/Linalg.h"
|
||||
#include "mlir/Dialect/SCF/IR/SCF.h"
|
||||
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
||||
#include "mlir/Pass/Pass.h"
|
||||
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
|
||||
#include "mlir/Transforms/DialectConversion.h"
|
||||
|
||||
#include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
|
||||
#include "mlir/Transforms/Passes.h"
|
||||
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
|
||||
#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/MatrixProductLowering.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp"
|
||||
#include "src/Accelerators/PIM/Pass/PIMPasses.h"
|
||||
|
||||
using namespace mlir;
|
||||
|
||||
namespace onnx_mlir {
|
||||
namespace {
|
||||
|
||||
static FailureOr<RowStripPhysicalValue> getRowStripValue(Value value) {
|
||||
return getRowStripPhysicalValue(value);
|
||||
}
|
||||
|
||||
static FailureOr<Value> publishRowStripValue(Operation* planOp,
|
||||
Value storage,
|
||||
PatternRewriter& rewriter) {
|
||||
auto logicalType = dyn_cast<RankedTensorType>(planOp->getResult(0).getType());
|
||||
if (!logicalType)
|
||||
return planOp->emitOpError("requires ranked logical output type"), failure();
|
||||
FailureOr<RowStripPhysicalValue> value = describeRowStripPhysicalValue(storage, logicalType);
|
||||
if (failed(value))
|
||||
return planOp->emitOpError("lowering produced invalid row-strip physical storage"), failure();
|
||||
FailureOr<Value> blueprint = createRowStripStorageBlueprint(
|
||||
storage, logicalType, rewriter, planOp->getLoc());
|
||||
if (failed(blueprint))
|
||||
return planOp->emitOpError("failed to create row-strip storage Blueprint"), failure();
|
||||
rewriter.replaceOp(planOp, *blueprint);
|
||||
return *blueprint;
|
||||
}
|
||||
|
||||
static bool isRowStripSelected(Operation* op) {
|
||||
auto selected = spatial::getSelectedPhysicalLayout(op);
|
||||
return selected && *selected == spatial::PhysicalLayout::NHWCRowStrip;
|
||||
}
|
||||
|
||||
static bool isDenseSelected(Operation* op) {
|
||||
auto selected = spatial::getSelectedPhysicalLayout(op);
|
||||
return selected && *selected == spatial::PhysicalLayout::DenseNCHW;
|
||||
}
|
||||
|
||||
static spatial::PhysicalLayout getKnownPhysicalLayout(Value value) {
|
||||
if (auto materialize = value.getDefiningOp<spatial::SpatMaterializeLayoutOp>())
|
||||
return materialize.getTargetPhysicalLayout();
|
||||
if (auto blueprint = value.getDefiningOp<spatial::SpatBlueprintOp>())
|
||||
return blueprint.getPhysicalLayout();
|
||||
if (Operation* producer = value.getDefiningOp()) {
|
||||
if (auto selected = spatial::getSelectedPhysicalLayout(producer))
|
||||
return *selected;
|
||||
}
|
||||
return spatial::PhysicalLayout::DenseNCHW;
|
||||
}
|
||||
|
||||
static LogicalResult verifySelectedLayouts(
|
||||
func::FuncOp funcOp, const spatial::SpatialTargetInfo& target) {
|
||||
LogicalResult result = success();
|
||||
funcOp.walk([&](Operation* op) {
|
||||
auto capability = dyn_cast<spatial::SpatialLayoutCapabilityInterface>(op);
|
||||
if (!capability)
|
||||
return;
|
||||
auto selected = spatial::getSelectedPhysicalLayout(op);
|
||||
if (!selected) {
|
||||
op->emitOpError("requires a selected physical layout from SpatialLayoutPlanning");
|
||||
result = failure();
|
||||
return;
|
||||
}
|
||||
if (*selected != spatial::PhysicalLayout::DenseNCHW
|
||||
&& *selected != spatial::PhysicalLayout::NHWCRowStrip) {
|
||||
op->emitOpError("has an unsupported selected physical layout");
|
||||
result = failure();
|
||||
return;
|
||||
}
|
||||
SmallVector<spatial::PhysicalLayout> operandLayouts;
|
||||
operandLayouts.reserve(op->getNumOperands());
|
||||
for (Value operand : op->getOperands())
|
||||
operandLayouts.push_back(getKnownPhysicalLayout(operand));
|
||||
auto alternatives = capability.getLayoutAlternatives(target, operandLayouts);
|
||||
if (llvm::none_of(alternatives, [&](const spatial::LayoutAlternative& alternative) {
|
||||
return alternative.resultLayout == *selected
|
||||
&& alternative.operandLayouts == operandLayouts;
|
||||
})) {
|
||||
op->emitOpError("selected physical layout is not lowerable for its explicit operand layouts");
|
||||
result = failure();
|
||||
}
|
||||
});
|
||||
return result;
|
||||
}
|
||||
|
||||
static FailureOr<Value>
|
||||
lowerRowStripRelu(const RowStripPhysicalValue& input, spatial::SpatReluPlanOp planOp, PatternRewriter& rewriter) {
|
||||
return applyRowStripRelu(input, rewriter, planOp.getLoc());
|
||||
}
|
||||
|
||||
static FailureOr<Value>
|
||||
lowerRowStripSilu(const RowStripPhysicalValue& input, spatial::SpatSiluPlanOp planOp, PatternRewriter& rewriter) {
|
||||
return applyRowStripSilu(input, rewriter, planOp.getLoc());
|
||||
}
|
||||
|
||||
static FailureOr<Value> lowerRowStripBiasAdd(const RowStripPhysicalValue& input,
|
||||
spatial::SpatBiasAddPlanOp planOp,
|
||||
PatternRewriter& rewriter) {
|
||||
return applyRowStripBiasAdd(input, planOp.getBias(), rewriter, planOp.getLoc());
|
||||
}
|
||||
|
||||
static FailureOr<Value> lowerRowStripAdd(const RowStripPhysicalValue& lhs,
|
||||
const RowStripPhysicalValue& rhs,
|
||||
spatial::SpatAddPlanOp planOp,
|
||||
PatternRewriter& rewriter) {
|
||||
return applyRowStripAdd(lhs, rhs, rewriter, planOp.getLoc());
|
||||
}
|
||||
|
||||
static FailureOr<Value> lowerRowStripConcat(ArrayRef<RowStripPhysicalValue> inputs,
|
||||
spatial::SpatConcatPlanOp planOp,
|
||||
PatternRewriter& rewriter) {
|
||||
auto outputType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
|
||||
if (!outputType)
|
||||
return failure();
|
||||
return applyRowStripConcat(inputs, outputType, rewriter, planOp.getLoc());
|
||||
}
|
||||
|
||||
static FailureOr<Value>
|
||||
materializeRowStripToDense(const RowStripPhysicalValue& rowStripValue, Location loc, PatternRewriter& rewriter) {
|
||||
if (rowStripValue.logicalType.getRank() != 4 || !rowStripValue.logicalType.hasStaticShape())
|
||||
return failure();
|
||||
return createRowStripAssemblyBlueprint(rowStripValue, rewriter, loc);
|
||||
}
|
||||
|
||||
static FailureOr<Value> materializeDenseToRowStrip(
|
||||
Value input, RankedTensorType logicalType, Location loc, PatternRewriter& rewriter) {
|
||||
if (!logicalType || !logicalType.hasStaticShape() || logicalType.getRank() != 4
|
||||
|| logicalType.getDimSize(0) != 1)
|
||||
return failure();
|
||||
auto nhwcType = RankedTensorType::get(
|
||||
{1, logicalType.getDimSize(2), logicalType.getDimSize(3), logicalType.getDimSize(1)},
|
||||
logicalType.getElementType(), logicalType.getEncoding());
|
||||
auto rowsType = RankedTensorType::get(
|
||||
{logicalType.getDimSize(2) * logicalType.getDimSize(3), logicalType.getDimSize(1)},
|
||||
logicalType.getElementType(), logicalType.getEncoding());
|
||||
auto rowsCompute = createSpatCompute<1>(
|
||||
rewriter, loc, rowsType, {}, input, [&](Value denseInput) {
|
||||
Value nhwc = createLinalgTranspose(
|
||||
denseInput, nhwcType, {0, 2, 3, 1}, rewriter, loc);
|
||||
Value rows = tensor::CollapseShapeOp::create(
|
||||
rewriter, loc, rowsType, nhwc,
|
||||
SmallVector<ReassociationIndices> {{0, 1, 2}, {3}});
|
||||
spatial::SpatYieldOp::create(rewriter, loc, rows);
|
||||
});
|
||||
Value rows = rowsCompute->getResult(0);
|
||||
FailureOr<Value> storage = createRowStripStorageFromRows(rows, logicalType, rewriter, loc);
|
||||
if (failed(storage))
|
||||
return failure();
|
||||
return createRowStripStorageBlueprint(*storage, logicalType, rewriter, loc);
|
||||
}
|
||||
|
||||
static FailureOr<Value> lowerDenseBatchBiasAdd(Value input, Value bias, RankedTensorType resultType,
|
||||
PatternRewriter& rewriter, Location loc) {
|
||||
auto producer = input.getDefiningOp<spatial::SpatGraphComputeBatch>();
|
||||
auto inputType = dyn_cast<RankedTensorType>(input.getType());
|
||||
auto biasType = dyn_cast<RankedTensorType>(bias.getType());
|
||||
if (!producer || !inputType || !biasType || !inputType.hasStaticShape() || !biasType.hasStaticShape()
|
||||
|| !resultType.hasStaticShape() || inputType.getDimSize(0) != producer.getLaneCount()
|
||||
|| biasType.getDimSize(0) != producer.getLaneCount() || resultType.getDimSize(0) != producer.getLaneCount())
|
||||
return failure();
|
||||
auto inputFragmentType = spatial::getGraphBatchFragmentType(inputType, producer.getLaneCount());
|
||||
auto outputFragmentType = spatial::getGraphBatchFragmentType(resultType, producer.getLaneCount());
|
||||
if (failed(inputFragmentType) || failed(outputFragmentType) || inputFragmentType->getRank() != biasType.getRank()
|
||||
|| inputFragmentType->getDimSize(0) != 1 || inputFragmentType->getShape().drop_front() != biasType.getShape().drop_front()
|
||||
|| inputFragmentType->getRank() != outputFragmentType->getRank() + 1)
|
||||
return failure();
|
||||
for (auto [inputDim, outputDim] : llvm::zip(inputFragmentType->getShape().drop_front(), outputFragmentType->getShape()))
|
||||
if (outputDim > inputDim)
|
||||
return failure();
|
||||
|
||||
auto batch = createSpatComputeBatch(rewriter, loc, TypeRange {resultType}, producer.getLaneCount(), {}, ValueRange {input, bias},
|
||||
[&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
|
||||
FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[0], args.lane, *inputFragmentType);
|
||||
if (failed(fragment))
|
||||
return failure();
|
||||
MixedSliceGeometry biasSlice;
|
||||
for (int64_t dim : inputFragmentType->getShape()) {
|
||||
biasSlice.offsets.push_back(biasSlice.offsets.empty() ? OpFoldResult(args.lane) : rewriter.getIndexAttr(0));
|
||||
biasSlice.sizes.push_back(rewriter.getIndexAttr(dim));
|
||||
biasSlice.strides.push_back(rewriter.getIndexAttr(1));
|
||||
}
|
||||
Value biasFragment = extractMixedSliceOrIdentity(rewriter, loc, args.inputs[1], *inputFragmentType, biasSlice);
|
||||
if (!biasFragment)
|
||||
return failure();
|
||||
Value added = spatial::SpatVAddOp::create(rewriter, loc, *inputFragmentType, *fragment, biasFragment);
|
||||
MixedSliceGeometry outputSlice;
|
||||
outputSlice.offsets.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(0));
|
||||
outputSlice.sizes.push_back(rewriter.getIndexAttr(1));
|
||||
outputSlice.strides.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(1));
|
||||
for (int64_t dim : outputFragmentType->getShape())
|
||||
outputSlice.sizes.push_back(rewriter.getIndexAttr(dim));
|
||||
Value output = extractMixedSliceOrIdentity(rewriter, loc, added, *outputFragmentType, outputSlice);
|
||||
if (!output)
|
||||
return failure();
|
||||
publishGraphBatchPhysicalFragment(rewriter, loc, output, args.outputs.front(), args.lane);
|
||||
return success();
|
||||
});
|
||||
if (failed(batch))
|
||||
return failure();
|
||||
return batch->getResult(0);
|
||||
}
|
||||
|
||||
struct LowerDenseReluPlan final : OpRewritePattern<spatial::SpatReluPlanOp> {
|
||||
using OpRewritePattern::OpRewritePattern;
|
||||
|
||||
LogicalResult matchAndRewrite(spatial::SpatReluPlanOp planOp,
|
||||
PatternRewriter& rewriter) const override {
|
||||
auto selected = spatial::getSelectedPhysicalLayout(planOp.getOperation());
|
||||
if (!selected || *selected != spatial::PhysicalLayout::DenseNCHW)
|
||||
return failure();
|
||||
|
||||
auto computeOp = createSpatCompute<1>(
|
||||
rewriter, planOp.getLoc(), planOp.getOutput().getType(), {}, planOp.getInput(), [&](Value x) {
|
||||
auto relu = spatial::SpatReluOp::create(rewriter, planOp.getLoc(), planOp.getOutput().getType(), x);
|
||||
spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), relu.getResult());
|
||||
});
|
||||
rewriter.replaceOp(planOp, computeOp.getResults());
|
||||
return success();
|
||||
}
|
||||
};
|
||||
|
||||
struct LowerDenseSiluPlan final : OpRewritePattern<spatial::SpatSiluPlanOp> {
|
||||
using OpRewritePattern::OpRewritePattern;
|
||||
|
||||
LogicalResult matchAndRewrite(spatial::SpatSiluPlanOp planOp,
|
||||
PatternRewriter& rewriter) const override {
|
||||
auto selected = spatial::getSelectedPhysicalLayout(planOp.getOperation());
|
||||
if (!selected || *selected != spatial::PhysicalLayout::DenseNCHW)
|
||||
return failure();
|
||||
|
||||
auto computeOp = createSpatCompute<1>(
|
||||
rewriter, planOp.getLoc(), planOp.getOutput().getType(), {}, planOp.getInput(), [&](Value x) {
|
||||
Value sigmoid = spatial::SpatSigmoidOp::create(
|
||||
rewriter, planOp.getLoc(), planOp.getOutput().getType(), x).getResult();
|
||||
Value silu = spatial::SpatVMulOp::create(
|
||||
rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, sigmoid).getResult();
|
||||
spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), silu);
|
||||
});
|
||||
rewriter.replaceOp(planOp, computeOp.getResults());
|
||||
return success();
|
||||
}
|
||||
};
|
||||
|
||||
struct LowerDenseResizePlan final : OpRewritePattern<spatial::SpatResizeNearestPlanOp> {
|
||||
explicit LowerDenseResizePlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
|
||||
: OpRewritePattern<spatial::SpatResizeNearestPlanOp>(ctx), target(target) {}
|
||||
|
||||
LogicalResult matchAndRewrite(spatial::SpatResizeNearestPlanOp planOp,
|
||||
PatternRewriter& rewriter) const override {
|
||||
if (!isDenseSelected(planOp.getOperation()))
|
||||
return failure();
|
||||
FailureOr<Value> lowered = lowerSelectedResizeNearestPlan(planOp, std::nullopt, target, rewriter);
|
||||
if (failed(lowered))
|
||||
return planOp.emitOpError("failed to lower selected dense nearest Resize plan");
|
||||
rewriter.replaceOp(planOp, *lowered);
|
||||
return success();
|
||||
}
|
||||
|
||||
const spatial::SpatialTargetInfo& target;
|
||||
};
|
||||
|
||||
struct LowerDenseBiasAddPlan final : OpRewritePattern<spatial::SpatBiasAddPlanOp> {
|
||||
using OpRewritePattern::OpRewritePattern;
|
||||
|
||||
LogicalResult matchAndRewrite(spatial::SpatBiasAddPlanOp planOp,
|
||||
PatternRewriter& rewriter) const override {
|
||||
if (!isDenseSelected(planOp.getOperation()))
|
||||
return failure();
|
||||
auto resultType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
|
||||
if (!resultType)
|
||||
return planOp.emitOpError("requires ranked output type");
|
||||
|
||||
FailureOr<Value> denseBias = materializeDenseBiasAddTensor(
|
||||
planOp.getBias(), resultType, rewriter, planOp.getLoc());
|
||||
if (failed(denseBias))
|
||||
return planOp.emitOpError("failed to materialize dense Conv-style bias");
|
||||
if (planOp.getInput().getDefiningOp<spatial::SpatGraphComputeBatch>()) {
|
||||
FailureOr<Value> lowered = lowerDenseBatchBiasAdd(
|
||||
planOp.getInput(), *denseBias, resultType, rewriter, planOp.getLoc());
|
||||
if (succeeded(lowered)) {
|
||||
rewriter.replaceOp(planOp, *lowered);
|
||||
return success();
|
||||
}
|
||||
}
|
||||
auto computeOp = createSpatCompute<2>(
|
||||
rewriter,
|
||||
planOp.getLoc(),
|
||||
planOp.getOutput().getType(),
|
||||
{},
|
||||
ValueRange {planOp.getInput(), *denseBias},
|
||||
[&](Value x, Value y) {
|
||||
auto added = spatial::SpatVAddOp::create(
|
||||
rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, y);
|
||||
spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added.getResult());
|
||||
});
|
||||
rewriter.replaceOp(planOp, computeOp.getResults());
|
||||
return success();
|
||||
}
|
||||
};
|
||||
|
||||
struct LowerDenseAddPlan final : OpRewritePattern<spatial::SpatAddPlanOp> {
|
||||
using OpRewritePattern::OpRewritePattern;
|
||||
|
||||
LogicalResult matchAndRewrite(spatial::SpatAddPlanOp planOp,
|
||||
PatternRewriter& rewriter) const override {
|
||||
if (!isDenseSelected(planOp.getOperation()))
|
||||
return failure();
|
||||
auto compute = createSpatCompute<2>(
|
||||
rewriter,
|
||||
planOp.getLoc(),
|
||||
planOp.getOutput().getType(),
|
||||
{},
|
||||
ValueRange {planOp.getLhs(), planOp.getRhs()},
|
||||
[&](Value lhsValue, Value rhsValue) {
|
||||
Value added = spatial::SpatVAddOp::create(
|
||||
rewriter, planOp.getLoc(), planOp.getOutput().getType(), lhsValue, rhsValue);
|
||||
spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added);
|
||||
});
|
||||
rewriter.replaceOp(planOp, compute.getResults());
|
||||
return success();
|
||||
}
|
||||
};
|
||||
|
||||
struct LowerDenseConcatPlan final : OpRewritePattern<spatial::SpatConcatPlanOp> {
|
||||
using OpRewritePattern::OpRewritePattern;
|
||||
|
||||
LogicalResult matchAndRewrite(spatial::SpatConcatPlanOp planOp,
|
||||
PatternRewriter& rewriter) const override {
|
||||
if (!isDenseSelected(planOp.getOperation()))
|
||||
return failure();
|
||||
auto compute = createSpatCompute(
|
||||
rewriter,
|
||||
planOp.getLoc(),
|
||||
TypeRange {planOp.getOutput().getType()},
|
||||
{},
|
||||
planOp.getInputs(),
|
||||
[&](ValueRange values) {
|
||||
Value concatenated = spatial::SpatConcatOp::create(
|
||||
rewriter,
|
||||
planOp.getLoc(),
|
||||
planOp.getOutput().getType(),
|
||||
rewriter.getI64IntegerAttr(planOp.getAxis()),
|
||||
values);
|
||||
spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), concatenated);
|
||||
});
|
||||
rewriter.replaceOp(planOp, compute.getResults());
|
||||
return success();
|
||||
}
|
||||
};
|
||||
|
||||
static LogicalResult lowerAddPlan(spatial::SpatAddPlanOp planOp,
|
||||
PatternRewriter& rewriter) {
|
||||
FailureOr<RowStripPhysicalValue> lhs = getRowStripValue(planOp.getLhs());
|
||||
FailureOr<RowStripPhysicalValue> rhs = getRowStripValue(planOp.getRhs());
|
||||
if (isRowStripSelected(planOp.getOperation()) && failed(lhs)) {
|
||||
if (getKnownPhysicalLayout(planOp.getLhs()) == spatial::PhysicalLayout::NHWCRowStrip)
|
||||
return failure();
|
||||
return planOp.emitOpError("selected row-strip Add plan requires row-strip inputs");
|
||||
}
|
||||
if (isRowStripSelected(planOp.getOperation()) && failed(rhs)) {
|
||||
if (getKnownPhysicalLayout(planOp.getRhs()) == spatial::PhysicalLayout::NHWCRowStrip)
|
||||
return failure();
|
||||
return planOp.emitOpError("selected row-strip Add plan requires row-strip inputs");
|
||||
}
|
||||
if (isRowStripSelected(planOp.getOperation())) {
|
||||
rewriter.setInsertionPoint(planOp);
|
||||
FailureOr<Value> lowered = lowerRowStripAdd(*lhs, *rhs, planOp, rewriter);
|
||||
if (failed(lowered))
|
||||
return planOp.emitOpError("failed to lower selected row-strip Spatial add plan");
|
||||
if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
|
||||
return failure();
|
||||
return success();
|
||||
}
|
||||
return planOp.emitOpError("dense Add plan was not lowered by the selected-plan patterns");
|
||||
}
|
||||
|
||||
static LogicalResult lowerConcatPlan(spatial::SpatConcatPlanOp planOp,
|
||||
PatternRewriter& rewriter) {
|
||||
SmallVector<RowStripPhysicalValue> inputs;
|
||||
for (Value input : planOp.getInputs()) {
|
||||
FailureOr<RowStripPhysicalValue> physical = getRowStripValue(input);
|
||||
if (failed(physical)) {
|
||||
inputs.clear();
|
||||
break;
|
||||
}
|
||||
inputs.push_back(*physical);
|
||||
}
|
||||
if (isRowStripSelected(planOp.getOperation()) && inputs.size() != planOp.getInputs().size()) {
|
||||
if (llvm::any_of(planOp.getInputs(), [](Value input) {
|
||||
return getKnownPhysicalLayout(input) == spatial::PhysicalLayout::NHWCRowStrip;
|
||||
}))
|
||||
return failure();
|
||||
return planOp.emitOpError("selected row-strip Concat plan requires row-strip inputs");
|
||||
}
|
||||
if (isRowStripSelected(planOp.getOperation())) {
|
||||
rewriter.setInsertionPoint(planOp);
|
||||
FailureOr<Value> lowered = lowerRowStripConcat(inputs, planOp, rewriter);
|
||||
if (failed(lowered))
|
||||
return planOp.emitOpError("failed to lower selected row-strip Spatial concat plan");
|
||||
if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
|
||||
return failure();
|
||||
return success();
|
||||
}
|
||||
return planOp.emitOpError("dense Concat plan was not lowered by the selected-plan patterns");
|
||||
}
|
||||
|
||||
struct LowerSelectedConvPlan final : OpRewritePattern<spatial::SpatConv2DPlanOp> {
|
||||
explicit LowerSelectedConvPlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
|
||||
: OpRewritePattern<spatial::SpatConv2DPlanOp>(ctx), target(target) {}
|
||||
|
||||
LogicalResult matchAndRewrite(spatial::SpatConv2DPlanOp planOp,
|
||||
PatternRewriter& rewriter) const override {
|
||||
if (isDenseSelected(planOp.getOperation())) {
|
||||
FailureOr<Value> lowered = lowerSelectedConv2DPlan(
|
||||
planOp, std::nullopt, /*emitRowStripLayout=*/false, target, rewriter);
|
||||
if (failed(lowered))
|
||||
return planOp.emitOpError("failed to lower selected dense Spatial Conv plan");
|
||||
rewriter.replaceOp(planOp, *lowered);
|
||||
return success();
|
||||
}
|
||||
if (!isRowStripSelected(planOp.getOperation()))
|
||||
return failure();
|
||||
|
||||
FailureOr<RowStripPhysicalValue> rowStripInput = getRowStripValue(planOp.getInput());
|
||||
if (failed(rowStripInput)
|
||||
&& getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
|
||||
return failure();
|
||||
std::optional<Value> physicalInput;
|
||||
if (succeeded(rowStripInput))
|
||||
physicalInput = rowStripInput->storage;
|
||||
FailureOr<Value> lowered = lowerSelectedConv2DPlan(
|
||||
planOp, physicalInput, /*emitRowStripLayout=*/true, target, rewriter);
|
||||
if (failed(lowered))
|
||||
return planOp.emitOpError("failed to lower selected row-strip Spatial Conv plan");
|
||||
if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
|
||||
return failure();
|
||||
return success();
|
||||
}
|
||||
|
||||
const spatial::SpatialTargetInfo& target;
|
||||
};
|
||||
|
||||
struct LowerRowStripReluPlan final : OpRewritePattern<spatial::SpatReluPlanOp> {
|
||||
using OpRewritePattern::OpRewritePattern;
|
||||
|
||||
LogicalResult matchAndRewrite(spatial::SpatReluPlanOp planOp,
|
||||
PatternRewriter& rewriter) const override {
|
||||
if (!isRowStripSelected(planOp.getOperation()))
|
||||
return failure();
|
||||
FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
|
||||
if (failed(input)) {
|
||||
if (getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
|
||||
return failure();
|
||||
return planOp.emitOpError("selected row-strip ReLU plan requires a row-strip input");
|
||||
}
|
||||
FailureOr<Value> lowered = lowerRowStripRelu(*input, planOp, rewriter);
|
||||
if (failed(lowered))
|
||||
return planOp.emitOpError("failed to lower selected row-strip Spatial ReLU plan");
|
||||
if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
|
||||
return failure();
|
||||
return success();
|
||||
}
|
||||
};
|
||||
|
||||
struct LowerRowStripSiluPlan final : OpRewritePattern<spatial::SpatSiluPlanOp> {
|
||||
using OpRewritePattern::OpRewritePattern;
|
||||
|
||||
LogicalResult matchAndRewrite(spatial::SpatSiluPlanOp planOp,
|
||||
PatternRewriter& rewriter) const override {
|
||||
if (!isRowStripSelected(planOp.getOperation()))
|
||||
return failure();
|
||||
FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
|
||||
if (failed(input)) {
|
||||
if (getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
|
||||
return failure();
|
||||
return planOp.emitOpError("selected row-strip SiLU plan requires a row-strip input");
|
||||
}
|
||||
FailureOr<Value> lowered = lowerRowStripSilu(*input, planOp, rewriter);
|
||||
if (failed(lowered))
|
||||
return planOp.emitOpError("failed to lower selected row-strip Spatial SiLU plan");
|
||||
if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
|
||||
return failure();
|
||||
return success();
|
||||
}
|
||||
};
|
||||
|
||||
struct LowerRowStripResizePlan final : OpRewritePattern<spatial::SpatResizeNearestPlanOp> {
|
||||
explicit LowerRowStripResizePlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
|
||||
: OpRewritePattern<spatial::SpatResizeNearestPlanOp>(ctx), target(target) {}
|
||||
|
||||
LogicalResult matchAndRewrite(spatial::SpatResizeNearestPlanOp planOp,
|
||||
PatternRewriter& rewriter) const override {
|
||||
if (!isRowStripSelected(planOp.getOperation()))
|
||||
return failure();
|
||||
FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
|
||||
if (failed(input)) {
|
||||
if (getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
|
||||
return failure();
|
||||
return planOp.emitOpError("selected row-strip Resize plan requires a row-strip input");
|
||||
}
|
||||
FailureOr<Value> lowered = lowerSelectedResizeNearestPlan(planOp, input->storage, target, rewriter);
|
||||
if (failed(lowered))
|
||||
return planOp.emitOpError("failed to lower selected row-strip Resize plan");
|
||||
if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
|
||||
return failure();
|
||||
return success();
|
||||
}
|
||||
|
||||
const spatial::SpatialTargetInfo& target;
|
||||
};
|
||||
|
||||
struct LowerDenseMaxPoolPlan final : OpRewritePattern<spatial::SpatMaxPool2DPlanOp> {
|
||||
explicit LowerDenseMaxPoolPlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
|
||||
: OpRewritePattern<spatial::SpatMaxPool2DPlanOp>(ctx), target(target) {}
|
||||
|
||||
LogicalResult matchAndRewrite(spatial::SpatMaxPool2DPlanOp planOp,
|
||||
PatternRewriter& rewriter) const override {
|
||||
if (!isDenseSelected(planOp.getOperation()))
|
||||
return failure();
|
||||
FailureOr<Value> lowered = lowerDenseMaxPool2DPlan(planOp, target, rewriter);
|
||||
if (failed(lowered))
|
||||
return planOp.emitOpError("failed to lower selected dense Spatial MaxPool plan");
|
||||
rewriter.replaceOp(planOp, *lowered);
|
||||
return success();
|
||||
}
|
||||
|
||||
const spatial::SpatialTargetInfo& target;
|
||||
};
|
||||
|
||||
struct LowerRowStripMaxPoolPlan final : OpRewritePattern<spatial::SpatMaxPool2DPlanOp> {
|
||||
explicit LowerRowStripMaxPoolPlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
|
||||
: OpRewritePattern<spatial::SpatMaxPool2DPlanOp>(ctx), target(target) {}
|
||||
|
||||
LogicalResult matchAndRewrite(spatial::SpatMaxPool2DPlanOp planOp,
|
||||
PatternRewriter& rewriter) const override {
|
||||
if (!isRowStripSelected(planOp.getOperation()))
|
||||
return failure();
|
||||
FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
|
||||
if (failed(input)
|
||||
&& getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
|
||||
return failure();
|
||||
std::optional<Value> physicalInput;
|
||||
if (succeeded(input))
|
||||
physicalInput = input->storage;
|
||||
FailureOr<Value> lowered = lowerSelectedMaxPool2DPlan(planOp, physicalInput, target, rewriter);
|
||||
if (failed(lowered))
|
||||
return planOp.emitOpError("failed to lower selected row-strip Spatial MaxPool plan");
|
||||
if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
|
||||
return failure();
|
||||
return success();
|
||||
}
|
||||
|
||||
const spatial::SpatialTargetInfo& target;
|
||||
};
|
||||
|
||||
struct LowerRowStripGlobalAveragePoolPlan
|
||||
final : OpRewritePattern<spatial::SpatGlobalAveragePoolPlanOp> {
|
||||
explicit LowerRowStripGlobalAveragePoolPlan(MLIRContext* ctx, const spatial::SpatialTargetInfo& target)
|
||||
: OpRewritePattern<spatial::SpatGlobalAveragePoolPlanOp>(ctx), target(target) {}
|
||||
|
||||
LogicalResult matchAndRewrite(spatial::SpatGlobalAveragePoolPlanOp planOp,
|
||||
PatternRewriter& rewriter) const override {
|
||||
if (!isRowStripSelected(planOp.getOperation()))
|
||||
return failure();
|
||||
FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
|
||||
if (failed(input)
|
||||
&& getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
|
||||
return failure();
|
||||
std::optional<Value> physicalInput;
|
||||
if (succeeded(input))
|
||||
physicalInput = input->storage;
|
||||
FailureOr<Value> lowered = lowerSelectedGlobalAveragePoolPlan(planOp, physicalInput, target, rewriter);
|
||||
if (failed(lowered))
|
||||
return planOp.emitOpError("failed to lower selected row-strip Spatial global AveragePool plan");
|
||||
if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
|
||||
return failure();
|
||||
return success();
|
||||
}
|
||||
|
||||
const spatial::SpatialTargetInfo& target;
|
||||
};
|
||||
|
||||
struct LowerDenseGlobalAveragePoolPlan
|
||||
final : OpRewritePattern<spatial::SpatGlobalAveragePoolPlanOp> {
|
||||
explicit LowerDenseGlobalAveragePoolPlan(MLIRContext* ctx,
|
||||
const spatial::SpatialTargetInfo& target)
|
||||
: OpRewritePattern<spatial::SpatGlobalAveragePoolPlanOp>(ctx), target(target) {}
|
||||
|
||||
LogicalResult matchAndRewrite(spatial::SpatGlobalAveragePoolPlanOp planOp,
|
||||
PatternRewriter& rewriter) const override {
|
||||
if (!isDenseSelected(planOp.getOperation()))
|
||||
return failure();
|
||||
FailureOr<Value> lowered = lowerDenseGlobalAveragePoolPlan(planOp, target, rewriter);
|
||||
if (failed(lowered))
|
||||
return planOp.emitOpError("failed to lower selected dense Spatial global AveragePool plan");
|
||||
rewriter.replaceOp(planOp, *lowered);
|
||||
return success();
|
||||
}
|
||||
|
||||
const spatial::SpatialTargetInfo& target;
|
||||
};
|
||||
|
||||
struct LowerRowStripBiasAddPlan final : OpRewritePattern<spatial::SpatBiasAddPlanOp> {
|
||||
using OpRewritePattern::OpRewritePattern;
|
||||
|
||||
LogicalResult matchAndRewrite(spatial::SpatBiasAddPlanOp planOp,
|
||||
PatternRewriter& rewriter) const override {
|
||||
if (!isRowStripSelected(planOp.getOperation()))
|
||||
return failure();
|
||||
FailureOr<RowStripPhysicalValue> input = getRowStripValue(planOp.getInput());
|
||||
if (failed(input)) {
|
||||
if (getKnownPhysicalLayout(planOp.getInput()) == spatial::PhysicalLayout::NHWCRowStrip)
|
||||
return failure();
|
||||
return planOp.emitOpError("selected row-strip bias_add plan requires a row-strip input");
|
||||
}
|
||||
FailureOr<Value> lowered = lowerRowStripBiasAdd(*input, planOp, rewriter);
|
||||
if (failed(lowered))
|
||||
return planOp.emitOpError("failed to lower selected row-strip Spatial bias_add plan");
|
||||
if (failed(publishRowStripValue(planOp, *lowered, rewriter)))
|
||||
return failure();
|
||||
return success();
|
||||
}
|
||||
};
|
||||
|
||||
struct LowerRowStripAddPlan final : OpRewritePattern<spatial::SpatAddPlanOp> {
|
||||
using OpRewritePattern::OpRewritePattern;
|
||||
|
||||
LogicalResult matchAndRewrite(spatial::SpatAddPlanOp planOp,
|
||||
PatternRewriter& rewriter) const override {
|
||||
if (!isRowStripSelected(planOp.getOperation()))
|
||||
return failure();
|
||||
return lowerAddPlan(planOp, rewriter);
|
||||
}
|
||||
};
|
||||
|
||||
struct LowerRowStripConcatPlan final : OpRewritePattern<spatial::SpatConcatPlanOp> {
|
||||
using OpRewritePattern::OpRewritePattern;
|
||||
|
||||
LogicalResult matchAndRewrite(spatial::SpatConcatPlanOp planOp,
|
||||
PatternRewriter& rewriter) const override {
|
||||
if (!isRowStripSelected(planOp.getOperation()))
|
||||
return failure();
|
||||
return lowerConcatPlan(planOp, rewriter);
|
||||
}
|
||||
};
|
||||
|
||||
struct LowerMaterializeLayout final
|
||||
: OpRewritePattern<spatial::SpatMaterializeLayoutOp> {
|
||||
using OpRewritePattern::OpRewritePattern;
|
||||
|
||||
LogicalResult matchAndRewrite(spatial::SpatMaterializeLayoutOp materializeOp,
|
||||
PatternRewriter& rewriter) const override {
|
||||
auto source = materializeOp.getSourcePhysicalLayout();
|
||||
auto target = materializeOp.getTargetPhysicalLayout();
|
||||
if (source == spatial::PhysicalLayout::DenseNCHW
|
||||
&& target == spatial::PhysicalLayout::DenseNCHW) {
|
||||
rewriter.replaceOp(materializeOp, materializeOp.getInput());
|
||||
return success();
|
||||
}
|
||||
if (source == spatial::PhysicalLayout::DenseNCHW
|
||||
&& target == spatial::PhysicalLayout::NHWCRowStrip) {
|
||||
auto logicalType = dyn_cast<RankedTensorType>(materializeOp.getInput().getType());
|
||||
if (!logicalType)
|
||||
return materializeOp.emitOpError("requires a ranked dense input"), failure();
|
||||
FailureOr<Value> rowStrip = materializeDenseToRowStrip(
|
||||
materializeOp.getInput(), logicalType, materializeOp.getLoc(), rewriter);
|
||||
if (failed(rowStrip))
|
||||
return materializeOp.emitOpError(
|
||||
"failed to materialize dense NCHW storage to row-strip layout"), failure();
|
||||
rewriter.replaceOp(materializeOp, *rowStrip);
|
||||
return success();
|
||||
}
|
||||
if (source != spatial::PhysicalLayout::NHWCRowStrip
|
||||
|| target != spatial::PhysicalLayout::DenseNCHW)
|
||||
return materializeOp.emitOpError(
|
||||
"unsupported Spatial layout materialization direction"), failure();
|
||||
auto inputType = dyn_cast<RankedTensorType>(materializeOp.getInput().getType());
|
||||
if (!inputType)
|
||||
return materializeOp.emitOpError("requires a ranked row-strip input"), failure();
|
||||
FailureOr<RowStripPhysicalValue> rowStripValue =
|
||||
getRowStripValue(materializeOp.getInput());
|
||||
if (failed(rowStripValue))
|
||||
return materializeOp.emitOpError(
|
||||
"requires an explicitly defining row-strip physical value"), failure();
|
||||
FailureOr<Value> dense = materializeRowStripToDense(
|
||||
*rowStripValue, materializeOp.getLoc(), rewriter);
|
||||
if (failed(dense))
|
||||
return materializeOp.emitOpError(
|
||||
"failed to materialize row-strip storage to dense NCHW"), failure();
|
||||
rewriter.replaceOp(materializeOp, *dense);
|
||||
return success();
|
||||
}
|
||||
};
|
||||
|
||||
struct LowerRowStripFlatten final
|
||||
: OpRewritePattern<spatial::SpatGraphCompute> {
|
||||
explicit LowerRowStripFlatten(MLIRContext* context,
|
||||
const spatial::SpatialTargetInfo& target)
|
||||
: OpRewritePattern<spatial::SpatGraphCompute>(context), target(target) {}
|
||||
|
||||
LogicalResult matchAndRewrite(spatial::SpatGraphCompute flattenOp,
|
||||
PatternRewriter& rewriter) const override {
|
||||
if (flattenOp.getInputs().size() != 1)
|
||||
return failure();
|
||||
FailureOr<RowStripPhysicalValue> input =
|
||||
getRowStripValue(flattenOp.getInputs().front());
|
||||
if (failed(input) || failed(canLowerFlattenFromRowStrip(flattenOp, target)))
|
||||
return failure();
|
||||
if (failed(lowerFlattenFromRowStrip(*input, flattenOp, target, rewriter)))
|
||||
return flattenOp.emitOpError(
|
||||
"failed to preserve row-strip layout through Flatten"), failure();
|
||||
return success();
|
||||
}
|
||||
|
||||
const spatial::SpatialTargetInfo& target;
|
||||
};
|
||||
|
||||
struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, OperationPass<ModuleOp>> {
|
||||
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerSpatialPlansPass)
|
||||
|
||||
StringRef getArgument() const override { return "lower-spatial-plans"; }
|
||||
StringRef getDescription() const override { return "Lower selected Spatial planning ops to low-level Spatial IR."; }
|
||||
|
||||
LowerSpatialPlansPass() = default;
|
||||
explicit LowerSpatialPlansPass(const spatial::SpatialTargetInfo& target)
|
||||
: target(target), hasTarget(true) {}
|
||||
|
||||
void runOnOperation() override {
|
||||
ModuleOp moduleOp = getOperation();
|
||||
if (!hasTarget) {
|
||||
moduleOp.emitError("Spatial plan lowering requires an injected SpatialTargetInfo");
|
||||
signalPassFailure();
|
||||
return;
|
||||
}
|
||||
MLIRContext* ctx = moduleOp.getContext();
|
||||
auto entryFunc = getPimEntryFunc(moduleOp);
|
||||
if (failed(entryFunc)) {
|
||||
moduleOp.emitError("failed to locate the PIM entry function during LowerSpatialPlans");
|
||||
signalPassFailure();
|
||||
return;
|
||||
}
|
||||
func::FuncOp funcOp = *entryFunc;
|
||||
PatternRewriter rewriter(ctx);
|
||||
auto verifyLogicalPhase = [&](StringRef stage) -> bool {
|
||||
if (succeeded(verifyLogicalSpatialGraphInvariants(*entryFunc)))
|
||||
return true;
|
||||
moduleOp.emitError() << "logical Spatial graph verification failed " << stage;
|
||||
signalPassFailure();
|
||||
return false;
|
||||
};
|
||||
|
||||
if (!verifyLogicalPhase("at the start of LowerSpatialPlans"))
|
||||
return;
|
||||
if (failed(verifySelectedLayouts(funcOp, target))) {
|
||||
moduleOp.emitError("selected Spatial layout verification failed");
|
||||
signalPassFailure();
|
||||
return;
|
||||
}
|
||||
|
||||
RewritePatternSet selectedPlanPatterns(ctx);
|
||||
selectedPlanPatterns.add<LowerDenseReluPlan,
|
||||
LowerRowStripReluPlan,
|
||||
LowerDenseSiluPlan,
|
||||
LowerRowStripSiluPlan,
|
||||
LowerDenseBiasAddPlan,
|
||||
LowerRowStripBiasAddPlan,
|
||||
LowerDenseAddPlan,
|
||||
LowerRowStripAddPlan,
|
||||
LowerDenseConcatPlan,
|
||||
LowerRowStripConcatPlan>(ctx);
|
||||
selectedPlanPatterns.add<LowerSelectedConvPlan,
|
||||
LowerDenseResizePlan,
|
||||
LowerRowStripResizePlan,
|
||||
LowerDenseMaxPoolPlan,
|
||||
LowerRowStripMaxPoolPlan,
|
||||
LowerDenseGlobalAveragePoolPlan,
|
||||
LowerRowStripGlobalAveragePoolPlan>(ctx, target);
|
||||
if (failed(applyPatternsGreedily(funcOp, std::move(selectedPlanPatterns)))) {
|
||||
moduleOp.emitError("failed to lower selected Spatial plans");
|
||||
signalPassFailure();
|
||||
return;
|
||||
}
|
||||
|
||||
RewritePatternSet layoutPatterns(ctx);
|
||||
layoutPatterns.add<LowerMaterializeLayout>(ctx);
|
||||
layoutPatterns.add<LowerRowStripFlatten>(ctx, target);
|
||||
ConversionTarget layoutTarget(*ctx);
|
||||
layoutTarget.addLegalDialect<spatial::SpatialDialect,
|
||||
tensor::TensorDialect,
|
||||
linalg::LinalgDialect,
|
||||
affine::AffineDialect,
|
||||
arith::ArithDialect,
|
||||
scf::SCFDialect,
|
||||
func::FuncDialect>();
|
||||
layoutTarget.addIllegalDialect<ONNXDialect>();
|
||||
layoutTarget.addIllegalOp<spatial::SpatMaterializeLayoutOp>();
|
||||
layoutTarget.addDynamicallyLegalOp<spatial::SpatGraphCompute>(
|
||||
[&](spatial::SpatGraphCompute computeOp) {
|
||||
if (computeOp.getInputs().size() != 1)
|
||||
return true;
|
||||
FailureOr<RowStripPhysicalValue> input =
|
||||
getRowStripValue(computeOp.getInputs().front());
|
||||
return failed(input) || failed(canLowerFlattenFromRowStrip(computeOp, target));
|
||||
});
|
||||
FrozenRewritePatternSet frozenLayoutPatterns(std::move(layoutPatterns));
|
||||
if (failed(applyFullConversion(funcOp, layoutTarget,
|
||||
frozenLayoutPatterns))) {
|
||||
moduleOp.emitError("failed to lower explicit Spatial layout materialization");
|
||||
signalPassFailure();
|
||||
return;
|
||||
}
|
||||
|
||||
if (!verifyLogicalPhase("after selected-plan conversion"))
|
||||
return;
|
||||
SmallVector<spatial::SpatBlueprintOp> deadPhysicalViews;
|
||||
funcOp.walk([&](spatial::SpatBlueprintOp blueprint) {
|
||||
if (spatial::isPhysicalView(blueprint.getMode()) && blueprint.use_empty())
|
||||
deadPhysicalViews.push_back(blueprint);
|
||||
});
|
||||
for (spatial::SpatBlueprintOp blueprint : deadPhysicalViews)
|
||||
rewriter.eraseOp(blueprint);
|
||||
bool hasIllegalOps = false;
|
||||
moduleOp.walk([&](Operation* op) {
|
||||
if (isa<ONNXEntryPointOp>(op))
|
||||
return;
|
||||
if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
|
||||
if (spatial::isFragmentAssembly(blueprint.getMode()))
|
||||
return;
|
||||
op->emitOpError("planning blueprint must not remain after LowerSpatialPlans");
|
||||
hasIllegalOps = true;
|
||||
}
|
||||
else if (isa<spatial::SpatConv2DPlanOp,
|
||||
spatial::SpatBiasAddPlanOp,
|
||||
spatial::SpatAddPlanOp,
|
||||
spatial::SpatReluPlanOp,
|
||||
spatial::SpatSiluPlanOp,
|
||||
spatial::SpatResizeNearestPlanOp,
|
||||
spatial::SpatMaxPool2DPlanOp,
|
||||
spatial::SpatGlobalAveragePoolPlanOp,
|
||||
spatial::SpatMaterializeLayoutOp>(op)
|
||||
|| op->getDialect()->getNamespace() == "onnx") {
|
||||
op->emitOpError("operation must not remain after LowerSpatialPlans");
|
||||
hasIllegalOps = true;
|
||||
}
|
||||
});
|
||||
|
||||
PassManager canonicalizationPM(ctx);
|
||||
canonicalizationPM.addPass(createCanonicalizerPass());
|
||||
if (failed(canonicalizationPM.run(moduleOp)))
|
||||
moduleOp.emitWarning("failed to run LowerSpatialPlansPass canonicalization; continuing");
|
||||
|
||||
if (hasIllegalOps) {
|
||||
signalPassFailure();
|
||||
} else {
|
||||
dumpModule(moduleOp, "spatial1_graph");
|
||||
spatial::SpatialDataflowExportStage exportMode = spatial::getSpatialDataflowExportStage();
|
||||
if (spatial::shouldExportSpatialDataflowStage(exportMode, spatial::SpatialDataflowExportStage::Spatial1)
|
||||
&& failed(spatial::exportSpatialDataflowCsvGraph(funcOp, "spatial1_graph"))) {
|
||||
signalPassFailure();
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
if (!verifyLogicalPhase("at the end of LowerSpatialPlans"))
|
||||
return;
|
||||
}
|
||||
|
||||
spatial::SpatialTargetInfo target;
|
||||
bool hasTarget = false;
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
std::unique_ptr<Pass> createLowerSpatialPlansPass() { return std::make_unique<LowerSpatialPlansPass>(); }
|
||||
|
||||
std::unique_ptr<Pass> createLowerSpatialPlansPass(const spatial::SpatialTargetInfo& target) {
|
||||
return std::make_unique<LowerSpatialPlansPass>(target);
|
||||
}
|
||||
|
||||
} // namespace onnx_mlir
|
||||
@@ -1,64 +0,0 @@
|
||||
#pragma once
|
||||
|
||||
#include <optional>
|
||||
|
||||
#include "mlir/IR/PatternMatch.h"
|
||||
#include "mlir/Support/LogicalResult.h"
|
||||
|
||||
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
|
||||
|
||||
namespace onnx_mlir {
|
||||
|
||||
struct RowStripPhysicalValue;
|
||||
|
||||
mlir::FailureOr<mlir::Value>
|
||||
lowerSelectedConv2DPlan(spatial::SpatConv2DPlanOp planOp,
|
||||
std::optional<mlir::Value> rowStripInput,
|
||||
bool emitRowStripLayout,
|
||||
const spatial::SpatialTargetInfo& target,
|
||||
mlir::PatternRewriter& rewriter);
|
||||
|
||||
mlir::LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp,
|
||||
const spatial::SpatialTargetInfo& target);
|
||||
mlir::LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp,
|
||||
const spatial::SpatialTargetInfo& target);
|
||||
|
||||
mlir::LogicalResult canLowerResizeNearestPlanToRowStrip(
|
||||
spatial::SpatResizeNearestPlanOp planOp, const spatial::SpatialTargetInfo& target);
|
||||
|
||||
mlir::FailureOr<mlir::Value> lowerSelectedResizeNearestPlan(
|
||||
spatial::SpatResizeNearestPlanOp planOp,
|
||||
std::optional<mlir::Value> rowStripInput,
|
||||
const spatial::SpatialTargetInfo& target,
|
||||
mlir::PatternRewriter& rewriter);
|
||||
|
||||
mlir::LogicalResult canLowerMaxPoolPlanToRowStrip(spatial::SpatMaxPool2DPlanOp planOp,
|
||||
const spatial::SpatialTargetInfo& target);
|
||||
|
||||
mlir::FailureOr<mlir::Value>
|
||||
lowerDenseMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
|
||||
const spatial::SpatialTargetInfo& target,
|
||||
mlir::PatternRewriter& rewriter);
|
||||
|
||||
mlir::FailureOr<mlir::Value>
|
||||
lowerSelectedMaxPool2DPlan(spatial::SpatMaxPool2DPlanOp planOp,
|
||||
std::optional<mlir::Value> rowStripInput,
|
||||
const spatial::SpatialTargetInfo& target,
|
||||
mlir::PatternRewriter& rewriter);
|
||||
|
||||
mlir::LogicalResult
|
||||
canLowerGlobalAveragePoolPlanToRowStrip(spatial::SpatGlobalAveragePoolPlanOp planOp,
|
||||
const spatial::SpatialTargetInfo& target);
|
||||
|
||||
mlir::FailureOr<mlir::Value>
|
||||
lowerDenseGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp,
|
||||
const spatial::SpatialTargetInfo& target,
|
||||
mlir::PatternRewriter& rewriter);
|
||||
|
||||
mlir::FailureOr<mlir::Value>
|
||||
lowerSelectedGlobalAveragePoolPlan(spatial::SpatGlobalAveragePoolPlanOp planOp,
|
||||
std::optional<mlir::Value> rowStripInput,
|
||||
const spatial::SpatialTargetInfo& target,
|
||||
mlir::PatternRewriter& rewriter);
|
||||
|
||||
} // namespace onnx_mlir
|
||||
@@ -1,133 +0,0 @@
|
||||
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/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 SpatialTargetInfo& 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> SpatReluPlanOp::getLayoutAlternatives(
|
||||
const SpatialTargetInfo&, 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 SpatialTargetInfo&, 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 SpatialTargetInfo& 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 SpatialTargetInfo& 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 SpatialTargetInfo& 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 SpatialTargetInfo&, 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 SpatialTargetInfo&, 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 SpatialTargetInfo&, 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
|
||||
@@ -1,265 +0,0 @@
|
||||
#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/Common/RowStripLayoutUtils.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
|
||||
#include "src/Accelerators/PIM/Pass/PIMPasses.h"
|
||||
|
||||
#include <algorithm>
|
||||
|
||||
using namespace mlir;
|
||||
|
||||
namespace onnx_mlir {
|
||||
namespace {
|
||||
|
||||
using LayoutMap = llvm::DenseMap<Value, spatial::PhysicalLayout>;
|
||||
|
||||
static spatial::PhysicalLayout getSelectedLayout(const LayoutMap& layouts, Value value) {
|
||||
if (auto it = layouts.find(value); it != layouts.end())
|
||||
return it->second;
|
||||
if (auto materialize = value.getDefiningOp<spatial::SpatMaterializeLayoutOp>())
|
||||
return materialize.getTargetPhysicalLayout();
|
||||
if (auto blueprint = value.getDefiningOp<spatial::SpatBlueprintOp>())
|
||||
return blueprint.getPhysicalLayout();
|
||||
return spatial::PhysicalLayout::DenseNCHW;
|
||||
}
|
||||
|
||||
static SmallVector<spatial::PhysicalLayout> getOperandLayouts(
|
||||
Operation* op, const LayoutMap& layouts) {
|
||||
SmallVector<spatial::PhysicalLayout> operandLayouts;
|
||||
operandLayouts.reserve(op->getNumOperands());
|
||||
for (Value operand : op->getOperands())
|
||||
operandLayouts.push_back(getSelectedLayout(layouts, operand));
|
||||
return operandLayouts;
|
||||
}
|
||||
|
||||
static FailureOr<SmallVector<spatial::LayoutAlternative>> getAlternatives(
|
||||
Operation* op, const LayoutMap& layouts, const spatial::SpatialTargetInfo& target) {
|
||||
auto capability = dyn_cast<spatial::SpatialLayoutCapabilityInterface>(op);
|
||||
if (!capability)
|
||||
return failure();
|
||||
SmallVector<spatial::LayoutAlternative> alternatives =
|
||||
capability.getLayoutAlternatives(target, getOperandLayouts(op, layouts));
|
||||
if (alternatives.empty())
|
||||
return op->emitOpError("does not advertise a legal Spatial layout alternative"), failure();
|
||||
for (const spatial::LayoutAlternative& alternative : alternatives)
|
||||
if (alternative.operandLayouts.size() != op->getNumOperands())
|
||||
return op->emitOpError("advertises a layout alternative with the wrong operand count"), failure();
|
||||
return alternatives;
|
||||
}
|
||||
|
||||
static unsigned findCurrentAlternative(
|
||||
Operation* op, ArrayRef<spatial::LayoutAlternative> alternatives,
|
||||
spatial::PhysicalLayout selectedResult) {
|
||||
for (auto [index, alternative] : llvm::enumerate(alternatives))
|
||||
if (alternative.resultLayout == selectedResult)
|
||||
return index;
|
||||
return 0;
|
||||
}
|
||||
|
||||
static int64_t alternativeCost(Operation* op,
|
||||
const spatial::LayoutAlternative& alternative,
|
||||
const LayoutMap& layouts,
|
||||
const LayoutMap& selectedResults,
|
||||
const spatial::SpatialTargetInfo& target) {
|
||||
int64_t cost = alternative.intrinsicCost;
|
||||
SmallVector<spatial::PhysicalLayout> operandLayouts = getOperandLayouts(op, layouts);
|
||||
for (auto [actual, required] : llvm::zip(operandLayouts, alternative.operandLayouts))
|
||||
cost += actual != required;
|
||||
|
||||
Value result = op->getResult(0);
|
||||
for (OpOperand& use : result.getUses()) {
|
||||
auto user = dyn_cast<spatial::SpatialLayoutCapabilityInterface>(use.getOwner());
|
||||
if (!user) {
|
||||
if (alternative.resultLayout != spatial::PhysicalLayout::DenseNCHW) {
|
||||
auto flatten = dyn_cast<spatial::SpatGraphCompute>(use.getOwner());
|
||||
if (!flatten || failed(canLowerFlattenFromRowStrip(flatten, target)))
|
||||
++cost;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
auto userAlternatives = getAlternatives(use.getOwner(), selectedResults, target);
|
||||
if (failed(userAlternatives))
|
||||
continue;
|
||||
spatial::PhysicalLayout userResult =
|
||||
selectedResults.lookup(use.getOwner()->getResult(0));
|
||||
unsigned userIndex = findCurrentAlternative(use.getOwner(), *userAlternatives, userResult);
|
||||
if (use.getOperandNumber() < (*userAlternatives)[userIndex].operandLayouts.size()
|
||||
&& (*userAlternatives)[userIndex].operandLayouts[use.getOperandNumber()]
|
||||
!= alternative.resultLayout)
|
||||
++cost;
|
||||
}
|
||||
return cost;
|
||||
}
|
||||
|
||||
static LogicalResult materializeMismatchedUses(
|
||||
IRRewriter& rewriter, Value value, const LayoutMap& layouts,
|
||||
const spatial::SpatialTargetInfo& target) {
|
||||
spatial::PhysicalLayout sourceLayout = getSelectedLayout(layouts, value);
|
||||
SmallVector<std::pair<OpOperand*, spatial::PhysicalLayout>> mismatches;
|
||||
for (OpOperand& use : value.getUses()) {
|
||||
Operation* userOp = use.getOwner();
|
||||
spatial::PhysicalLayout required = spatial::PhysicalLayout::DenseNCHW;
|
||||
if (auto capability = dyn_cast<spatial::SpatialLayoutCapabilityInterface>(userOp)) {
|
||||
auto alternatives = getAlternatives(userOp, layouts, target);
|
||||
if (failed(alternatives))
|
||||
return failure();
|
||||
spatial::PhysicalLayout selected =
|
||||
getSelectedLayout(layouts, userOp->getResult(0));
|
||||
unsigned selectedIndex = findCurrentAlternative(userOp, *alternatives, selected);
|
||||
required = (*alternatives)[selectedIndex].operandLayouts[use.getOperandNumber()];
|
||||
}
|
||||
else if (auto flatten = dyn_cast<spatial::SpatGraphCompute>(userOp);
|
||||
flatten && sourceLayout == spatial::PhysicalLayout::NHWCRowStrip
|
||||
&& succeeded(canLowerFlattenFromRowStrip(flatten, target))) {
|
||||
continue;
|
||||
}
|
||||
if (required != sourceLayout)
|
||||
mismatches.push_back({&use, required});
|
||||
}
|
||||
|
||||
for (auto [use, required] : mismatches) {
|
||||
Operation* userOp = use->getOwner();
|
||||
rewriter.setInsertionPoint(userOp);
|
||||
auto materialized = spatial::SpatMaterializeLayoutOp::create(
|
||||
rewriter, userOp->getLoc(), use->get().getType(), use->get(),
|
||||
spatial::LogicalLayoutAttr::get(
|
||||
rewriter.getContext(), spatial::LogicalLayout::NCHW),
|
||||
spatial::PhysicalLayoutAttr::get(rewriter.getContext(), sourceLayout),
|
||||
spatial::PhysicalLayoutAttr::get(rewriter.getContext(),
|
||||
required));
|
||||
use->set(materialized.getResult());
|
||||
}
|
||||
return success();
|
||||
}
|
||||
|
||||
static LogicalResult verifySelectedLayouts(
|
||||
ArrayRef<Operation*> planOps, const LayoutMap& layouts,
|
||||
const spatial::SpatialTargetInfo& target) {
|
||||
for (Operation* op : planOps) {
|
||||
auto selected = spatial::getSelectedPhysicalLayout(op);
|
||||
if (!selected)
|
||||
return op->emitOpError("requires a selected physical layout"), failure();
|
||||
auto alternatives = getAlternatives(op, layouts, target);
|
||||
if (failed(alternatives))
|
||||
return failure();
|
||||
if (llvm::none_of(*alternatives, [&](const spatial::LayoutAlternative& alternative) {
|
||||
return alternative.resultLayout == *selected;
|
||||
}))
|
||||
return op->emitOpError("selected physical layout is not advertised by its layout contract"), failure();
|
||||
}
|
||||
return success();
|
||||
}
|
||||
|
||||
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 Spatial layout alternatives and insert explicit reconciliation barriers.";
|
||||
}
|
||||
|
||||
SpatialLayoutPlanningPass() = default;
|
||||
explicit SpatialLayoutPlanningPass(const spatial::SpatialTargetInfo& target)
|
||||
: target(target), hasTarget(true) {}
|
||||
|
||||
void runOnOperation() override {
|
||||
ModuleOp moduleOp = getOperation();
|
||||
if (!hasTarget) {
|
||||
moduleOp.emitError("Spatial layout planning requires an injected SpatialTargetInfo");
|
||||
signalPassFailure();
|
||||
return;
|
||||
}
|
||||
auto entryFunc = getPimEntryFunc(moduleOp);
|
||||
if (failed(entryFunc)) {
|
||||
moduleOp.emitError("failed to locate the PIM entry function during Spatial layout planning");
|
||||
signalPassFailure();
|
||||
return;
|
||||
}
|
||||
|
||||
func::FuncOp funcOp = *entryFunc;
|
||||
SmallVector<Operation*> planOps;
|
||||
for (Operation& op : funcOp.getBody().front())
|
||||
if (isa<spatial::SpatialLayoutCapabilityInterface>(&op))
|
||||
planOps.push_back(&op);
|
||||
|
||||
LayoutMap layouts;
|
||||
for (Operation* op : planOps)
|
||||
layouts[op->getResult(0)] = spatial::PhysicalLayout::DenseNCHW;
|
||||
|
||||
const size_t maxRounds = 2 * planOps.size() + 1;
|
||||
bool converged = false;
|
||||
for (size_t round = 0; round < maxRounds && !converged; ++round) {
|
||||
converged = true;
|
||||
SmallVector<Operation*> order(planOps);
|
||||
if (round % 2)
|
||||
std::reverse(order.begin(), order.end());
|
||||
for (Operation* op : order) {
|
||||
auto alternatives = getAlternatives(op, layouts, target);
|
||||
if (failed(alternatives)) {
|
||||
signalPassFailure();
|
||||
return;
|
||||
}
|
||||
spatial::PhysicalLayout current = layouts.lookup(op->getResult(0));
|
||||
unsigned currentIndex = findCurrentAlternative(op, *alternatives, current);
|
||||
int64_t bestCost = alternativeCost(
|
||||
op, (*alternatives)[currentIndex], layouts, layouts, target);
|
||||
unsigned bestIndex = currentIndex;
|
||||
for (auto [index, alternative] : llvm::enumerate(*alternatives)) {
|
||||
int64_t cost = alternativeCost(op, alternative, layouts, layouts, target);
|
||||
if (cost < bestCost) {
|
||||
bestCost = cost;
|
||||
bestIndex = index;
|
||||
}
|
||||
}
|
||||
spatial::PhysicalLayout selected = (*alternatives)[bestIndex].resultLayout;
|
||||
if (selected != current) {
|
||||
layouts[op->getResult(0)] = selected;
|
||||
converged = false;
|
||||
}
|
||||
}
|
||||
}
|
||||
if (!converged) {
|
||||
moduleOp.emitError("Spatial layout selection did not converge within its bounded iteration budget");
|
||||
signalPassFailure();
|
||||
return;
|
||||
}
|
||||
IRRewriter rewriter(&getContext());
|
||||
for (Operation* op : planOps) {
|
||||
op->setAttr(spatial::kSelectedLayoutAttrName,
|
||||
spatial::PhysicalLayoutAttr::get(
|
||||
rewriter.getContext(), layouts.lookup(op->getResult(0))));
|
||||
if (failed(materializeMismatchedUses(rewriter, op->getResult(0), layouts, target))) {
|
||||
signalPassFailure();
|
||||
return;
|
||||
}
|
||||
}
|
||||
if (failed(verifySelectedLayouts(planOps, layouts, target))
|
||||
|| failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) {
|
||||
moduleOp.emitError("Spatial layout planning verification failed");
|
||||
signalPassFailure();
|
||||
}
|
||||
}
|
||||
|
||||
spatial::SpatialTargetInfo target;
|
||||
bool hasTarget = false;
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
std::unique_ptr<Pass> createSpatialLayoutPlanningPass() {
|
||||
return std::make_unique<SpatialLayoutPlanningPass>();
|
||||
}
|
||||
|
||||
std::unique_ptr<Pass> createSpatialLayoutPlanningPass(
|
||||
const spatial::SpatialTargetInfo& target) {
|
||||
return std::make_unique<SpatialLayoutPlanningPass>(target);
|
||||
}
|
||||
|
||||
} // namespace onnx_mlir
|
||||
-16
@@ -1,16 +0,0 @@
|
||||
#pragma once
|
||||
|
||||
#include "ScheduledComputeMaterialization.hpp"
|
||||
#include "Scheduling/MergeSchedulingAnalysis.hpp"
|
||||
|
||||
#include <memory>
|
||||
#include <optional>
|
||||
|
||||
namespace onnx_mlir::spatial {
|
||||
|
||||
struct ScheduledSpatialState {
|
||||
std::optional<MergeScheduleResult> logicalSchedule;
|
||||
std::optional<ScheduledComputeMaterializationResult> materialization;
|
||||
};
|
||||
|
||||
} // namespace onnx_mlir::spatial
|
||||
@@ -1,24 +0,0 @@
|
||||
#ifndef SPATIAL_LAYOUT_INTERFACE_TD
|
||||
#define SPATIAL_LAYOUT_INTERFACE_TD
|
||||
|
||||
include "mlir/IR/OpBase.td"
|
||||
|
||||
def SpatialLayoutCapabilityInterface : OpInterface<"SpatialLayoutCapabilityInterface"> {
|
||||
let description = [{
|
||||
Contract implemented by logical Spatial planning operations that expose
|
||||
their legal physical layout alternatives to the Spatial planner.
|
||||
}];
|
||||
|
||||
let methods = [
|
||||
InterfaceMethod<
|
||||
"Return legal physical layout alternatives for this operation and its current operand layouts.",
|
||||
"::llvm::SmallVector<::onnx_mlir::spatial::LayoutAlternative>",
|
||||
"getLayoutAlternatives",
|
||||
(ins "const ::onnx_mlir::spatial::SpatialTargetInfo &":$target,
|
||||
"::llvm::ArrayRef<::onnx_mlir::spatial::PhysicalLayout>":$operandLayouts)>
|
||||
];
|
||||
|
||||
let cppNamespace = "::onnx_mlir::spatial";
|
||||
}
|
||||
|
||||
#endif
|
||||
@@ -1,37 +0,0 @@
|
||||
#pragma once
|
||||
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
|
||||
namespace onnx_mlir::spatial {
|
||||
|
||||
struct MatrixUnitShape {
|
||||
size_t rows = 128;
|
||||
size_t columns = 128;
|
||||
};
|
||||
|
||||
enum class ConvLoweringStrategy : uint8_t {
|
||||
Auto,
|
||||
Legacy,
|
||||
Depthwise,
|
||||
PackedIm2Col,
|
||||
StreamedPatch,
|
||||
StreamedPacked,
|
||||
OutputChannelTiled,
|
||||
InputKTiled,
|
||||
Tiled2D,
|
||||
};
|
||||
|
||||
struct SpatialTargetInfo {
|
||||
MatrixUnitShape matrixShape;
|
||||
size_t matrixUnitsPerProcessor = 64;
|
||||
size_t processorCount = 1;
|
||||
size_t vectorWidth = 16;
|
||||
|
||||
uint64_t convIm2colMaxElements = 1ull << 20;
|
||||
uint64_t convStreamChunkPositions = 1024;
|
||||
ConvLoweringStrategy convLoweringStrategy = ConvLoweringStrategy::Auto;
|
||||
bool useExperimentalConvImplementation = false;
|
||||
};
|
||||
|
||||
} // namespace onnx_mlir::spatial
|
||||
@@ -1,63 +0,0 @@
|
||||
#pragma once
|
||||
|
||||
#include "mlir/Pass/Pass.h"
|
||||
|
||||
#include <cstddef>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
|
||||
namespace onnx_mlir {
|
||||
namespace spatial {
|
||||
struct SchedulingTarget;
|
||||
struct ScheduledSpatialState;
|
||||
struct SpatialTargetInfo;
|
||||
|
||||
std::unique_ptr<mlir::Pass> createScheduleSpatialGraphPass();
|
||||
std::unique_ptr<mlir::Pass> createScheduleSpatialGraphPass(const SchedulingTarget& target);
|
||||
std::unique_ptr<mlir::Pass> createScheduleSpatialGraphPass(
|
||||
const SchedulingTarget& target,
|
||||
std::shared_ptr<ScheduledSpatialState> state);
|
||||
std::unique_ptr<mlir::Pass> createVerifyScheduledSpatialPass();
|
||||
std::unique_ptr<mlir::Pass> createVerifyScheduledSpatialPass(
|
||||
std::shared_ptr<ScheduledSpatialState> state);
|
||||
std::unique_ptr<mlir::Pass> createRealizeSpatialCommunicationPass();
|
||||
std::unique_ptr<mlir::Pass> createRealizeSpatialCommunicationPass(
|
||||
const SchedulingTarget& target,
|
||||
std::shared_ptr<ScheduledSpatialState> state);
|
||||
std::unique_ptr<mlir::Pass> createVerifyRealizedSpatialPass();
|
||||
std::unique_ptr<mlir::Pass> createVerifyRealizedSpatialPass(
|
||||
std::shared_ptr<ScheduledSpatialState> state);
|
||||
}
|
||||
|
||||
std::unique_ptr<mlir::Pass> createONNXToSpatialPass();
|
||||
std::unique_ptr<mlir::Pass> createONNXToSpatialPass(const spatial::SpatialTargetInfo& target);
|
||||
std::unique_ptr<mlir::Pass> createSpatialLayoutPlanningPass();
|
||||
std::unique_ptr<mlir::Pass> createSpatialLayoutPlanningPass(const spatial::SpatialTargetInfo& target);
|
||||
std::unique_ptr<mlir::Pass> createLowerSpatialPlansPass();
|
||||
std::unique_ptr<mlir::Pass> createLowerSpatialPlansPass(const spatial::SpatialTargetInfo& target);
|
||||
|
||||
std::unique_ptr<mlir::Pass> createSpatialToPimPass();
|
||||
|
||||
std::unique_ptr<mlir::Pass> createPimBufferizationPreparationPass();
|
||||
std::unique_ptr<mlir::Pass> createPimOneShotBufferizationPass();
|
||||
std::unique_ptr<mlir::Pass> createPimMemoryNormalizationPass();
|
||||
std::unique_ptr<mlir::Pass> createPimBufferizationVerificationPass();
|
||||
|
||||
|
||||
std::unique_ptr<mlir::Pass> createTrivialGraphComputeMergePass();
|
||||
std::unique_ptr<mlir::Pass> createTrivialGraphComputeMergePass(
|
||||
size_t residentWeightCapacity);
|
||||
|
||||
std::unique_ptr<mlir::Pass> createPimHostConstantFoldingPass();
|
||||
|
||||
std::unique_ptr<mlir::Pass> createPimInstructionSelectionPass();
|
||||
|
||||
std::unique_ptr<mlir::Pass> createPimLocalMemoryPlanningPass();
|
||||
|
||||
std::unique_ptr<mlir::Pass> createPimVerificationPass();
|
||||
|
||||
std::unique_ptr<mlir::Pass> createEmitPimCodePass();
|
||||
|
||||
std::unique_ptr<mlir::Pass> createMessagePass(std::string message);
|
||||
|
||||
} // namespace onnx_mlir
|
||||
Reference in New Issue
Block a user