Raptor sync wait
This commit is contained in:
ilgeco
@@ -32,6 +32,9 @@ inline constexpr llvm::StringLiteral kCoreIdAttrName = "coreId";
|
||||
inline constexpr llvm::StringLiteral kCoreIdsAttrName = "coreIds";
|
||||
inline constexpr llvm::StringLiteral kLocalMemoryAddressAttrName = "pim.local_memory_address";
|
||||
inline constexpr llvm::StringLiteral kLocalMemorySizeAttrName = "pim.local_memory_size";
|
||||
inline constexpr llvm::StringLiteral kPipelineHostBufferBytesAttrName = "pim.pipeline_host_buffer_bytes";
|
||||
inline constexpr llvm::StringLiteral kPipelineHostBufferName = "pim_pipeline_channels";
|
||||
inline constexpr size_t kPimEventRegisterCount = 32;
|
||||
inline constexpr std::array<llvm::StringLiteral, 4> kRemovedLocalMemoryPlanAttrNames = {
|
||||
"pim.local_memory_slot",
|
||||
"pim.local_memory_slot_size",
|
||||
|
||||
@@ -162,8 +162,8 @@ inline constexpr std::array<InstructionJsonFormat, kOpcodeCount> kInstructionJso
|
||||
{true, true, true, "", "", "", "len" }, // lmv
|
||||
{true, false, true, "core", "", "", "size"}, // send
|
||||
{true, false, true, "core", "", "", "size"}, // recv
|
||||
{false, false, false, "", "", "", "" }, // wait
|
||||
{false, false, false, "", "", "", "" }, // sync
|
||||
{false, false, false, "", "event_register", "wait_value", ""}, // wait
|
||||
{false, false, false, "core", "event_register", "", ""}, // sync
|
||||
}};
|
||||
static_assert(kInstructionJsonFormats.size() == kOpcodeCount);
|
||||
|
||||
|
||||
@@ -692,6 +692,34 @@ void PimCodeGen::codeGenSendOp(pim::PimSendOp sendOp, const StaticValueKnowledge
|
||||
pim_binary::Opcode::send, addressOf(sendOp.getInput(), knowledge), *targetCoreId, sendOp.getSize());
|
||||
}
|
||||
|
||||
void PimCodeGen::codeGenWaitOp(
|
||||
pim::PimWaitOp waitOp, const StaticValueKnowledge& knowledge) const {
|
||||
auto eventRegister = indexOf(waitOp.getEventRegister(), knowledge);
|
||||
assert(succeeded(eventRegister)
|
||||
&& "pim.wait event register must be statically resolvable during codegen");
|
||||
pim_binary::InstructionRecord instruction;
|
||||
instruction.opcode = pim_binary::Opcode::wait;
|
||||
instruction.generic1 = pim::checkedI32OrCrash(
|
||||
*eventRegister, "wait event register");
|
||||
instruction.generic2 = waitOp.getWaitValue();
|
||||
emitInstruction(instruction);
|
||||
}
|
||||
|
||||
void PimCodeGen::codeGenSyncOp(
|
||||
pim::PimSyncOp syncOp, const StaticValueKnowledge& knowledge) const {
|
||||
auto targetCoreId = indexOf(syncOp.getTargetCoreId(), knowledge);
|
||||
auto eventRegister = indexOf(syncOp.getEventRegister(), knowledge);
|
||||
assert(succeeded(targetCoreId) && succeeded(eventRegister)
|
||||
&& "pim.sync operands must be statically resolvable during codegen");
|
||||
pim_binary::InstructionRecord instruction;
|
||||
instruction.opcode = pim_binary::Opcode::sync;
|
||||
instruction.r2OrImm = pim::checkedI32OrCrash(
|
||||
*targetCoreId, "sync target core id");
|
||||
instruction.generic1 = pim::checkedI32OrCrash(
|
||||
*eventRegister, "sync event register");
|
||||
emitInstruction(instruction);
|
||||
}
|
||||
|
||||
void PimCodeGen::codeGenConcatOp(pim::PimConcatOp concatOp, const StaticValueKnowledge& knowledge) const {
|
||||
auto outputType = cast<ShapedType>(concatOp.getOutputBuffer().getType());
|
||||
assert(outputType.hasStaticShape() && "concat codegen requires static output shape");
|
||||
@@ -991,6 +1019,8 @@ static LogicalResult executeCompiledCorePlan(
|
||||
case CompiledCoreOpKind::VMV: coreCodeGen.codeGenVMVOp(cast<pim::PimVMVOp>(node.op), knowledge); break;
|
||||
case CompiledCoreOpKind::Receive: coreCodeGen.codeGenReceiveOp(cast<pim::PimReceiveOp>(node.op), knowledge); break;
|
||||
case CompiledCoreOpKind::Send: coreCodeGen.codeGenSendOp(cast<pim::PimSendOp>(node.op), knowledge); break;
|
||||
case CompiledCoreOpKind::Wait: coreCodeGen.codeGenWaitOp(cast<pim::PimWaitOp>(node.op), knowledge); break;
|
||||
case CompiledCoreOpKind::Sync: coreCodeGen.codeGenSyncOp(cast<pim::PimSyncOp>(node.op), knowledge); break;
|
||||
case CompiledCoreOpKind::Concat: coreCodeGen.codeGenConcatOp(cast<pim::PimConcatOp>(node.op), knowledge); break;
|
||||
case CompiledCoreOpKind::Vmm:
|
||||
if (auto weightSlot = resolveWeightSlot(cast<pim::PimVMMOp>(node.op), knowledge); succeeded(weightSlot))
|
||||
|
||||
@@ -217,6 +217,8 @@ public:
|
||||
|
||||
void codeGenReceiveOp(pim::PimReceiveOp receiveOp, const StaticValueKnowledge& knowledge) const;
|
||||
void codeGenSendOp(pim::PimSendOp sendOp, const StaticValueKnowledge& knowledge) const;
|
||||
void codeGenWaitOp(pim::PimWaitOp waitOp, const StaticValueKnowledge& knowledge) const;
|
||||
void codeGenSyncOp(pim::PimSyncOp syncOp, const StaticValueKnowledge& knowledge) const;
|
||||
void codeGenConcatOp(pim::PimConcatOp concatOp, const StaticValueKnowledge& knowledge) const;
|
||||
|
||||
template <typename MVMTy>
|
||||
|
||||
@@ -2,6 +2,8 @@
|
||||
|
||||
#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
|
||||
|
||||
#include <limits>
|
||||
|
||||
#define DEBUG_TYPE "PimCompilerOptions"
|
||||
|
||||
namespace onnx_mlir {
|
||||
@@ -110,6 +112,12 @@ llvm::cl::opt<size_t>
|
||||
llvm::cl::opt<size_t>
|
||||
crossbarCountInCore("crossbar-count", llvm::cl::desc("Number of crossbars in each core"), llvm::cl::init(64));
|
||||
|
||||
llvm::cl::opt<size_t> pipelineStages(
|
||||
"pipeline",
|
||||
llvm::cl::desc("Number of throughput pipeline stages (1 preserves latency scheduling)"),
|
||||
llvm::cl::init(1),
|
||||
llvm::cl::cat(OnnxMlirOptions));
|
||||
|
||||
llvm::cl::opt<long> coresCount("core-count",
|
||||
llvm::cl::desc("Number of cores in the chip. Required for PIM compilation."),
|
||||
llvm::cl::init(-1));
|
||||
@@ -129,4 +137,14 @@ void verifyExplicitPimCoreCount() {
|
||||
llvm::report_fatal_error("PIM compilation requires --core-count to be a positive integer");
|
||||
}
|
||||
|
||||
void verifyPimPipelineStages() {
|
||||
if (pipelineStages.getValue() == 0)
|
||||
llvm::report_fatal_error("PIM compilation requires --pipeline to be positive");
|
||||
if (static_cast<size_t>(coresCount.getValue()) % pipelineStages.getValue() != 0)
|
||||
llvm::report_fatal_error("PIM compilation requires --core-count to be divisible by --pipeline");
|
||||
if (crossbarCountInCore.getValue()
|
||||
> std::numeric_limits<size_t>::max() / pipelineStages.getValue())
|
||||
llvm::report_fatal_error("PIM compilation --crossbar-count * --pipeline overflows");
|
||||
}
|
||||
|
||||
} // namespace onnx_mlir
|
||||
|
||||
@@ -62,6 +62,7 @@ extern llvm::cl::opt<bool> pimVerifyBufferizationCopyFreedom;
|
||||
|
||||
extern llvm::cl::opt<size_t> crossbarSize;
|
||||
extern llvm::cl::opt<size_t> crossbarCountInCore;
|
||||
extern llvm::cl::opt<size_t> pipelineStages;
|
||||
extern llvm::cl::opt<long> coresCount;
|
||||
extern llvm::cl::opt<std::string> pimTargetConfig;
|
||||
extern llvm::cl::opt<uint64_t> pimConvIm2colMaxElements;
|
||||
@@ -69,5 +70,6 @@ extern llvm::cl::opt<uint64_t> pimConvStreamChunkPositions;
|
||||
|
||||
bool hasExplicitPimCoreCount();
|
||||
void verifyExplicitPimCoreCount();
|
||||
void verifyPimPipelineStages();
|
||||
|
||||
} // namespace onnx_mlir
|
||||
|
||||
@@ -330,6 +330,7 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
|
||||
EmissionTargetType& emissionTarget,
|
||||
std::string outputNameNoExt) {
|
||||
verifyExplicitPimCoreCount();
|
||||
verifyPimPipelineStages();
|
||||
spatial::SchedulingTarget schedulingTarget = getPimSchedulingTarget();
|
||||
spatial::SpatialTargetResources targetResources = getPimSpatialTargetResources(schedulingTarget);
|
||||
|
||||
@@ -354,7 +355,7 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
|
||||
pm.addPass(createTrivialGraphComputeMergePass(
|
||||
schedulingTarget.residentWeightCapacity, exportStage));
|
||||
pm.addPass(spatial::createScheduleAndRealizeSpatialPass(
|
||||
schedulingTarget, exportStage));
|
||||
schedulingTarget, exportStage, pipelineStages.getValue()));
|
||||
pm.addPass(createMessagePass("Onnx lowered to Spatial"));
|
||||
}
|
||||
|
||||
|
||||
@@ -17,6 +17,8 @@ static FailureOr<CompiledCoreOpKind> classifyCompiledCoreOpKind(Operation& op) {
|
||||
if (isa<pim::PimVMVOp>(op)) return CompiledCoreOpKind::VMV;
|
||||
if (isa<pim::PimReceiveOp>(op)) return CompiledCoreOpKind::Receive;
|
||||
if (isa<pim::PimSendOp>(op)) return CompiledCoreOpKind::Send;
|
||||
if (isa<pim::PimWaitOp>(op)) return CompiledCoreOpKind::Wait;
|
||||
if (isa<pim::PimSyncOp>(op)) return CompiledCoreOpKind::Sync;
|
||||
if (isa<pim::PimConcatOp>(op)) return CompiledCoreOpKind::Concat;
|
||||
if (isa<pim::PimVMMOp>(op)) return CompiledCoreOpKind::Vmm;
|
||||
if (isa<pim::PimVVAddOp>(op)) return CompiledCoreOpKind::VVAdd;
|
||||
|
||||
@@ -17,6 +17,8 @@ enum class CompiledCoreOpKind : uint8_t {
|
||||
VMV,
|
||||
Receive,
|
||||
Send,
|
||||
Wait,
|
||||
Sync,
|
||||
Concat,
|
||||
Vmm,
|
||||
VVAdd,
|
||||
|
||||
@@ -0,0 +1,39 @@
|
||||
#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
|
||||
@@ -0,0 +1,23 @@
|
||||
#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
|
||||
@@ -0,0 +1,902 @@
|
||||
#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
|
||||
@@ -108,7 +108,9 @@ void verifyScheduledInputs(ComputeOpTy compute,
|
||||
for (auto [inputIndex, input] : llvm::enumerate(compute.getInputs())) {
|
||||
size_t currentInputIndex = inputIndex;
|
||||
Operation* definingOp = input.getDefiningOp();
|
||||
if (allowChannelReceiveInputs && isa_and_nonnull<spatial::SpatChannelReceiveOp>(definingOp))
|
||||
if (allowChannelReceiveInputs
|
||||
&& isa_and_nonnull<spatial::SpatChannelReceiveOp,
|
||||
spatial::SpatHostWaitLoadOp>(definingOp))
|
||||
continue;
|
||||
if (isScheduledPhase1Value(input))
|
||||
continue;
|
||||
@@ -163,7 +165,8 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter
|
||||
});
|
||||
continue;
|
||||
}
|
||||
if (isa<spatial::SpatChannelReceiveOp, spatial::SpatChannelSendOp>(&op)) {
|
||||
if (isa<spatial::SpatChannelReceiveOp, spatial::SpatChannelSendOp,
|
||||
spatial::SpatHostStoreSyncOp, spatial::SpatHostWaitLoadOp>(&op)) {
|
||||
diagnostics.report(&op, [&](Operation* illegalOp) {
|
||||
illegalOp->emitOpError() << kPhaseMarker
|
||||
<< " explicit channel communication is not expected before merge materialization";
|
||||
@@ -182,7 +185,8 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter
|
||||
|
||||
void verifyScheduledTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter& diagnostics) {
|
||||
for (Operation& op : funcOp.getOps()) {
|
||||
if (isa<spatial::SpatChannelSendOp, spatial::SpatChannelReceiveOp>(&op)) {
|
||||
if (isa<spatial::SpatChannelSendOp, spatial::SpatChannelReceiveOp,
|
||||
spatial::SpatHostStoreSyncOp, spatial::SpatHostWaitLoadOp>(&op)) {
|
||||
diagnostics.report(&op, [&](Operation* illegalOp) {
|
||||
illegalOp->emitOpError() << kPhaseMarker << " real channel communication is not allowed in scheduled phase 1";
|
||||
});
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
#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
|
||||
@@ -0,0 +1,133 @@
|
||||
#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
|
||||
@@ -0,0 +1,265 @@
|
||||
#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
|
||||
@@ -1,7 +1,10 @@
|
||||
#include "mlir/IR/ValueRange.h"
|
||||
|
||||
#include "mlir/Dialect/Arith/IR/Arith.h"
|
||||
#include "mlir/Dialect/Func/IR/FuncOps.h"
|
||||
#include "mlir/Dialect/MemRef/IR/MemRef.h"
|
||||
#include "mlir/Dialect/SCF/IR/SCF.h"
|
||||
#include "mlir/IR/BuiltinOps.h"
|
||||
|
||||
#include "llvm/ADT/STLExtras.h"
|
||||
|
||||
@@ -28,6 +31,49 @@ FailureOr<IntegerAttr> getTensorSizeInBytesAttr(Builder& builder, Operation* anc
|
||||
return pim::getCheckedI32Attr(builder, anchor, *byteSize, "tensor byte size");
|
||||
}
|
||||
|
||||
LogicalResult materializePipelineHostBuffer(
|
||||
func::FuncOp funcOp, RewriterBase &rewriter) {
|
||||
auto bytes = funcOp->getAttrOfType<IntegerAttr>(
|
||||
kPipelineHostBufferBytesAttrName);
|
||||
if (!bytes)
|
||||
return success();
|
||||
if (bytes.getInt() <= 0)
|
||||
return funcOp.emitOpError(
|
||||
"pipeline host transfer buffer must be positive");
|
||||
ModuleOp moduleOp = funcOp->getParentOfType<ModuleOp>();
|
||||
if (moduleOp.lookupSymbol<memref::GlobalOp>(kPipelineHostBufferName))
|
||||
return funcOp.emitOpError(
|
||||
"pipeline host transfer buffer symbol already exists");
|
||||
auto type = MemRefType::get(
|
||||
{bytes.getInt()}, rewriter.getI8Type());
|
||||
OpBuilder::InsertionGuard guard(rewriter);
|
||||
rewriter.setInsertionPointToStart(moduleOp.getBody());
|
||||
memref::GlobalOp::create(
|
||||
rewriter, funcOp.getLoc(),
|
||||
rewriter.getStringAttr(kPipelineHostBufferName),
|
||||
rewriter.getStringAttr("private"), TypeAttr::get(type), Attribute(),
|
||||
UnitAttr(), IntegerAttr());
|
||||
return success();
|
||||
}
|
||||
|
||||
FailureOr<mlir::Value> getPipelineHostBuffer(
|
||||
OpBuilder &builder, Operation *anchor) {
|
||||
auto funcOp = anchor->getParentOfType<func::FuncOp>();
|
||||
auto moduleOp = anchor->getParentOfType<ModuleOp>();
|
||||
auto bytes = funcOp
|
||||
? funcOp->getAttrOfType<IntegerAttr>(kPipelineHostBufferBytesAttrName)
|
||||
: IntegerAttr();
|
||||
auto global = moduleOp
|
||||
? moduleOp.lookupSymbol<memref::GlobalOp>(kPipelineHostBufferName)
|
||||
: memref::GlobalOp();
|
||||
if (!bytes || !global)
|
||||
return anchor->emitOpError(
|
||||
"requires the pipeline host transfer buffer"), failure();
|
||||
auto type = MemRefType::get({bytes.getInt()}, builder.getI8Type());
|
||||
return memref::GetGlobalOp::create(
|
||||
builder, anchor->getLoc(), type, kPipelineHostBufferName).getResult();
|
||||
}
|
||||
|
||||
Operation* getEarliestUserWithinBlock(mlir::Value value) {
|
||||
auto users = value.getUsers();
|
||||
|
||||
|
||||
@@ -10,6 +10,7 @@
|
||||
#include "mlir/IR/Builders.h"
|
||||
#include "mlir/IR/Value.h"
|
||||
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
||||
#include "mlir/Dialect/Func/IR/FuncOps.h"
|
||||
#include "mlir/Support/LogicalResult.h"
|
||||
|
||||
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
|
||||
@@ -23,6 +24,12 @@ namespace onnx_mlir {
|
||||
mlir::FailureOr<mlir::IntegerAttr>
|
||||
getTensorSizeInBytesAttr(mlir::Builder& builder, mlir::Operation* anchor, mlir::Value value);
|
||||
|
||||
mlir::LogicalResult materializePipelineHostBuffer(
|
||||
mlir::func::FuncOp funcOp, mlir::RewriterBase &rewriter);
|
||||
|
||||
mlir::FailureOr<mlir::Value> getPipelineHostBuffer(
|
||||
mlir::OpBuilder &builder, mlir::Operation *anchor);
|
||||
|
||||
template <class T>
|
||||
size_t rangeLength(const mlir::iterator_range<T> range) {
|
||||
return std::distance(range.begin(), range.end());
|
||||
|
||||
@@ -345,20 +345,39 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatScheduledCom
|
||||
auto blockArg = computeOp.getInputArgument(inputIndex);
|
||||
if (!blockArg)
|
||||
return computeOp.emitOpError("expected compute input block arguments during lowering");
|
||||
auto receiveOp = dyn_cast_or_null<spatial::SpatChannelReceiveOp>(input.getDefiningOp());
|
||||
auto channelReceive = dyn_cast_or_null<spatial::SpatChannelReceiveOp>(
|
||||
input.getDefiningOp());
|
||||
auto hostWaitLoad = dyn_cast_or_null<spatial::SpatHostWaitLoadOp>(
|
||||
input.getDefiningOp());
|
||||
Operation *receiveOp = channelReceive
|
||||
? channelReceive.getOperation() : hostWaitLoad.getOperation();
|
||||
if (receiveOp && !blockArg->use_empty()) {
|
||||
rewriter.setInsertionPoint(getEarliestUserWithinBlock(*blockArg));
|
||||
auto outputType = cast<ShapedType>(blockArg->getType());
|
||||
auto outputBuffer = createEmptyTensorFromShaped(rewriter, receiveOp.getLoc(), outputType);
|
||||
auto outputBuffer = createEmptyTensorFromShaped(
|
||||
rewriter, receiveOp->getLoc(), outputType);
|
||||
auto sizeAttr = getTensorSizeInBytesAttr(rewriter, computeOp.getOperation(), *blockArg);
|
||||
if (failed(sizeAttr))
|
||||
return failure();
|
||||
Value received =
|
||||
PimReceiveOp::create(
|
||||
rewriter, receiveOp.getLoc(), outputBuffer.getType(), outputBuffer,
|
||||
arith::ConstantIndexOp::create(rewriter, receiveOp.getLoc(), 0),
|
||||
*sizeAttr, receiveOp.getSourceCoreId())
|
||||
Value zero = arith::ConstantIndexOp::create(
|
||||
rewriter, receiveOp->getLoc(), 0);
|
||||
Value received;
|
||||
if (hostWaitLoad) {
|
||||
auto hostBuffer = getPipelineHostBuffer(rewriter, hostWaitLoad);
|
||||
if (failed(hostBuffer))
|
||||
return failure();
|
||||
PimWaitOp::create(
|
||||
rewriter, receiveOp->getLoc(), hostWaitLoad.getEventRegister(),
|
||||
rewriter.getI32IntegerAttr(1));
|
||||
received = PimMemCopyHostToDevOp::create(
|
||||
rewriter, receiveOp->getLoc(), outputBuffer.getType(), zero,
|
||||
hostWaitLoad.getHostOffset(), outputBuffer, *hostBuffer, *sizeAttr)
|
||||
.getOutput();
|
||||
} else {
|
||||
received = PimReceiveOp::create(
|
||||
rewriter, receiveOp->getLoc(), outputBuffer.getType(), outputBuffer,
|
||||
zero, *sizeAttr, channelReceive.getSourceCoreId()).getOutput();
|
||||
}
|
||||
blockArg->replaceAllUsesWith(received);
|
||||
markOpToRemove(receiveOp);
|
||||
continue;
|
||||
@@ -383,7 +402,8 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatScheduledCom
|
||||
if (rangeLength(resultUses) == 1) {
|
||||
OpOperand& resultUse = *resultUses.begin();
|
||||
Operation* resultUser = resultUse.getOwner();
|
||||
if (isa<spatial::SpatChannelSendOp>(resultUser))
|
||||
if (isa<spatial::SpatChannelSendOp,
|
||||
spatial::SpatHostStoreSyncOp>(resultUser))
|
||||
continue;
|
||||
}
|
||||
|
||||
|
||||
@@ -57,10 +57,29 @@ struct ChannelSendLowering : OpRewritePattern<spatial::SpatChannelSendOp> {
|
||||
}
|
||||
};
|
||||
|
||||
struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp> {
|
||||
struct HostStoreSyncLowering : OpRewritePattern<spatial::SpatHostStoreSyncOp> {
|
||||
using OpRewritePattern::OpRewritePattern;
|
||||
|
||||
LogicalResult matchAndRewrite(spatial::SpatChannelReceiveOp op, PatternRewriter& rewriter) const override {
|
||||
LogicalResult matchAndRewrite(spatial::SpatHostStoreSyncOp op, PatternRewriter& rewriter) const override {
|
||||
auto sizeAttr = getTensorSizeInBytesAttr(rewriter, op.getOperation(), op.getInput());
|
||||
auto hostBuffer = getPipelineHostBuffer(rewriter, op);
|
||||
if (failed(sizeAttr) || failed(hostBuffer))
|
||||
return failure();
|
||||
Value zero = arith::ConstantIndexOp::create(rewriter, op.getLoc(), 0);
|
||||
pim::PimMemCopyDevToHostOp::create(
|
||||
rewriter, op.getLoc(), hostBuffer->getType(), op.getHostOffset(), zero,
|
||||
*hostBuffer, op.getInput(), *sizeAttr);
|
||||
auto sync = pim::PimSyncOp::create(
|
||||
rewriter, op.getLoc(), op.getTargetCoreId(), op.getEventRegister());
|
||||
copyRaptorDebugAttrs(op.getOperation(), sync.getOperation());
|
||||
rewriter.eraseOp(op);
|
||||
return success();
|
||||
}
|
||||
};
|
||||
|
||||
template <typename ReceiveOp, typename CreateReceive>
|
||||
static LogicalResult lowerReceive(
|
||||
ReceiveOp op, PatternRewriter& rewriter, CreateReceive createReceive) {
|
||||
if (op->use_empty()) {
|
||||
rewriter.eraseOp(op);
|
||||
return success();
|
||||
@@ -86,12 +105,11 @@ struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp>
|
||||
if (failed(sizeAttr))
|
||||
return failure();
|
||||
Value zero = arith::ConstantIndexOp::create(rewriter, op.getLoc(), 0);
|
||||
auto receive = pim::PimReceiveOp::create(
|
||||
rewriter, op.getLoc(), op.getResult().getType(), outputBuffer, zero, *sizeAttr, op.getSourceCoreId());
|
||||
copyRaptorDebugAttrs(op.getOperation(), receive.getOperation());
|
||||
Value received = receive.getOutput();
|
||||
auto received = createReceive(outputBuffer, zero, *sizeAttr);
|
||||
if (failed(received))
|
||||
return failure();
|
||||
if (!destinationInsert) {
|
||||
rewriter.replaceOp(op, received);
|
||||
rewriter.replaceOp(op, *received);
|
||||
return success();
|
||||
}
|
||||
|
||||
@@ -99,10 +117,42 @@ struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp>
|
||||
Value targetOffset = createDestinationByteOffset(rewriter, destinationInsert);
|
||||
auto copy = pim::PimMemCopyOp::create(
|
||||
rewriter, op.getLoc(), destinationInsert.getDestType(), targetOffset, zero,
|
||||
destinationInsert.getDest(), received, *sizeAttr);
|
||||
destinationInsert.getDest(), *received, *sizeAttr);
|
||||
rewriter.replaceOp(destinationInsert, copy.getOutput());
|
||||
rewriter.eraseOp(op);
|
||||
return success();
|
||||
}
|
||||
|
||||
struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp> {
|
||||
using OpRewritePattern::OpRewritePattern;
|
||||
|
||||
LogicalResult matchAndRewrite(spatial::SpatChannelReceiveOp op, PatternRewriter& rewriter) const override {
|
||||
return lowerReceive(op, rewriter, [&](Value outputBuffer, Value zero, IntegerAttr sizeAttr) -> FailureOr<Value> {
|
||||
auto receive = pim::PimReceiveOp::create(
|
||||
rewriter, op.getLoc(), op.getResult().getType(), outputBuffer, zero,
|
||||
sizeAttr, op.getSourceCoreId());
|
||||
copyRaptorDebugAttrs(op.getOperation(), receive.getOperation());
|
||||
return receive.getOutput();
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
struct HostWaitLoadLowering : OpRewritePattern<spatial::SpatHostWaitLoadOp> {
|
||||
using OpRewritePattern::OpRewritePattern;
|
||||
|
||||
LogicalResult matchAndRewrite(spatial::SpatHostWaitLoadOp op, PatternRewriter& rewriter) const override {
|
||||
return lowerReceive(op, rewriter, [&](Value outputBuffer, Value zero, IntegerAttr sizeAttr) -> FailureOr<Value> {
|
||||
auto hostBuffer = getPipelineHostBuffer(rewriter, op);
|
||||
if (failed(hostBuffer))
|
||||
return failure();
|
||||
auto wait = pim::PimWaitOp::create(
|
||||
rewriter, op.getLoc(), op.getEventRegister(),
|
||||
rewriter.getI32IntegerAttr(1));
|
||||
copyRaptorDebugAttrs(op.getOperation(), wait.getOperation());
|
||||
return pim::PimMemCopyHostToDevOp::create(
|
||||
rewriter, op.getLoc(), outputBuffer.getType(), zero,
|
||||
op.getHostOffset(), outputBuffer, *hostBuffer, sizeAttr).getOutput();
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
@@ -148,7 +198,9 @@ struct ConcatLowering : OpRewritePattern<spatial::SpatConcatOp> {
|
||||
} // namespace
|
||||
|
||||
void populateChannelLoweringPatterns(RewritePatternSet& patterns) {
|
||||
patterns.add<ChannelSendLowering, ChannelReceiveLowering, ExtractRowsLowering, ConcatLowering>(patterns.getContext());
|
||||
patterns.add<ChannelSendLowering, ChannelReceiveLowering,
|
||||
HostStoreSyncLowering, HostWaitLoadLowering,
|
||||
ExtractRowsLowering, ConcatLowering>(patterns.getContext());
|
||||
}
|
||||
|
||||
} // namespace onnx_mlir
|
||||
|
||||
@@ -859,6 +859,10 @@ void raptor::SpatialToPimPass::replaceReturnWithOutputBuffers(func::ReturnOp ret
|
||||
markOpToRemove(receiveOp);
|
||||
return;
|
||||
}
|
||||
if (auto receiveOp = dyn_cast<spatial::SpatHostWaitLoadOp>(op)) {
|
||||
markOpToRemove(receiveOp);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
SmallVector<Value> originalOperands(returnOp.getOperands().begin(), returnOp.getOperands().end());
|
||||
|
||||
@@ -126,6 +126,8 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
|
||||
spatial::SpatConcatOp,
|
||||
spatial::SpatChannelReceiveOp,
|
||||
spatial::SpatChannelSendOp,
|
||||
spatial::SpatHostStoreSyncOp,
|
||||
spatial::SpatHostWaitLoadOp,
|
||||
spatial::SpatExtractRowsOp>();
|
||||
|
||||
RewritePatternSet initialPatterns(ctx);
|
||||
@@ -140,6 +142,12 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
|
||||
populateGlobalTensorMaterializationPatterns(globalTensorPatterns);
|
||||
walkAndApplyPatterns(moduleOp, std::move(globalTensorPatterns));
|
||||
|
||||
if (funcOp->hasAttr(kPipelineHostBufferBytesAttrName)
|
||||
&& failed(materializePipelineHostBuffer(funcOp, rewriter))) {
|
||||
signalPassFailure();
|
||||
return;
|
||||
}
|
||||
|
||||
auto returnOp = cast<func::ReturnOp>(funcOp.front().getTerminator());
|
||||
addReturnOutputBuffers(returnOp, rewriter);
|
||||
if (failed(allocateAndInitializeCoreLocalVariables(funcOp, rewriter))) {
|
||||
@@ -182,6 +190,17 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
SmallVector<spatial::SpatHostWaitLoadOp> hostWaitLoadOps;
|
||||
for (auto op : funcOp.getOps<spatial::SpatHostWaitLoadOp>())
|
||||
hostWaitLoadOps.push_back(op);
|
||||
for (auto op : hostWaitLoadOps) {
|
||||
bool onlyPendingRemovalUsers = llvm::all_of(
|
||||
op->getUsers(), [&](Operation* user) {
|
||||
return llvm::is_contained(operationsToRemove, user);
|
||||
});
|
||||
if (onlyPendingRemovalUsers)
|
||||
markOpToRemove(op);
|
||||
}
|
||||
|
||||
RewritePatternSet coreBodyPatterns(ctx);
|
||||
populateCoreBodyPatterns(coreBodyPatterns);
|
||||
@@ -202,6 +221,8 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
|
||||
spatial::SpatConcatOp,
|
||||
spatial::SpatChannelReceiveOp,
|
||||
spatial::SpatChannelSendOp,
|
||||
spatial::SpatHostStoreSyncOp,
|
||||
spatial::SpatHostWaitLoadOp,
|
||||
spatial::SpatExtractRowsOp>();
|
||||
|
||||
SmallVector<pim::PimCoreOp> coreOps;
|
||||
@@ -251,6 +272,8 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
|
||||
communicationTarget.addIllegalOp<spatial::SpatConcatOp,
|
||||
spatial::SpatChannelReceiveOp,
|
||||
spatial::SpatChannelSendOp,
|
||||
spatial::SpatHostStoreSyncOp,
|
||||
spatial::SpatHostWaitLoadOp,
|
||||
spatial::SpatExtractRowsOp>();
|
||||
|
||||
RewritePatternSet communicationPatterns(ctx);
|
||||
|
||||
@@ -430,8 +430,7 @@ analyzeCopyRewrite(Value target, Value source, Value targetOffset, Value sourceO
|
||||
|
||||
auto targetBytes = getShapedByteSize(targetType);
|
||||
auto sourceBytes = getShapedByteSize(sourceType);
|
||||
if (targetType.getElementType() == sourceType.getElementType() && succeeded(targetBytes) && succeeded(sourceBytes)
|
||||
&& size <= *targetBytes && size <= *sourceBytes) {
|
||||
if (succeeded(targetBytes) && succeeded(sourceBytes) && size <= *targetBytes && size <= *sourceBytes) {
|
||||
auto targetSuffixRank = getContiguousSuffixRank(target, targetType.getShape());
|
||||
auto sourceSuffixRank = getContiguousSuffixRank(source, sourceType.getShape());
|
||||
if (succeeded(targetSuffixRank) && succeeded(sourceSuffixRank)
|
||||
|
||||
@@ -241,6 +241,8 @@ static bool isSupportedCoreInstructionOp(Operation* op) {
|
||||
pim::PimVMVOp,
|
||||
pim::PimReceiveOp,
|
||||
pim::PimSendOp,
|
||||
pim::PimSyncOp,
|
||||
pim::PimWaitOp,
|
||||
pim::PimConcatOp,
|
||||
pim::PimVMMOp,
|
||||
pim::PimVVAddOp,
|
||||
|
||||
@@ -118,6 +118,32 @@ def PimReceiveOp : PimOp<"receive", [DestinationStyleOpInterface]> {
|
||||
}];
|
||||
}
|
||||
|
||||
def PimSyncOp : PimOp<"sync", []> {
|
||||
let summary = "Signal an event register on another core";
|
||||
|
||||
let arguments = (ins
|
||||
Index:$targetCoreId,
|
||||
Index:$eventRegister
|
||||
);
|
||||
|
||||
let assemblyFormat = [{
|
||||
$targetCoreId `event` $eventRegister attr-dict
|
||||
}];
|
||||
}
|
||||
|
||||
def PimWaitOp : PimOp<"wait", []> {
|
||||
let summary = "Wait for an event register value";
|
||||
|
||||
let arguments = (ins
|
||||
Index:$eventRegister,
|
||||
I32Attr:$waitValue
|
||||
);
|
||||
|
||||
let assemblyFormat = [{
|
||||
$eventRegister `value` $waitValue attr-dict
|
||||
}];
|
||||
}
|
||||
|
||||
def PimMemCopyHostToDevOp : PimOp<"memcp_hd", [DestinationStyleOpInterface]> {
|
||||
let summary = "Copy a memory region from host memory into device memory";
|
||||
|
||||
|
||||
@@ -35,6 +35,7 @@ add_pim_library(SpatialOps
|
||||
Passes/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp
|
||||
Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp
|
||||
Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp
|
||||
Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.cpp
|
||||
Passes/Transforms/TrivialGraphComputeMergePass.cpp
|
||||
|
||||
EXCLUDE_FROM_OM_LIBS
|
||||
|
||||
+4
-2
@@ -219,10 +219,12 @@ static void appendReceive(BoundaryProgram &boundary,
|
||||
run->entryOffsets[run->entryOffsets.size() - 2]].family->requirement;
|
||||
CollectionTarget previousTarget {run->collection, run->positions.back()};
|
||||
bool sameEntry = previous == requirement;
|
||||
if (sameEntry
|
||||
bool sameRoute = run->slices.back().family->hostRouted
|
||||
== slice.family->hostRouted;
|
||||
if (sameRoute && (sameEntry
|
||||
|| (sameCollectionEmissionContract(previousTarget, target)
|
||||
&& previous->publicationFragmentType
|
||||
== requirement->publicationFragmentType)) {
|
||||
== requirement->publicationFragmentType))) {
|
||||
run->slices.push_back(slice);
|
||||
if (sameEntry) {
|
||||
run->entryOffsets.back() = run->slices.size();
|
||||
|
||||
+130
-18
@@ -8,6 +8,7 @@
|
||||
#include "src/Accelerators/PIM/Common/IR/StaticIntGrid.hpp"
|
||||
#include "src/Accelerators/PIM/Common/IR/StaticIntSequence.hpp"
|
||||
#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
|
||||
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
|
||||
#include <array>
|
||||
namespace onnx_mlir::spatial {
|
||||
using namespace mlir;
|
||||
@@ -18,6 +19,8 @@ struct LogicalTransferMetadataView {
|
||||
StaticIntSequenceChain parentCounts;
|
||||
StaticIntSequenceChain sourceCores;
|
||||
StaticIntSequenceChain targetCores;
|
||||
StaticIntSequenceChain hostOffsets;
|
||||
StaticIntSequenceChain eventRegisters;
|
||||
StaticIntSequenceChain targetLanes;
|
||||
StaticIntSequenceChain localOffsets;
|
||||
SmallVector<StaticIntSequenceChain> projectionOffsets;
|
||||
@@ -28,7 +31,8 @@ struct LogicalTransferMetadataView {
|
||||
};
|
||||
using MetadataMember = StaticIntSequenceChain LogicalTransferMetadataView::*;
|
||||
static constexpr std::array<MetadataMember, 3> transferMetadataMembers{
|
||||
&LogicalTransferMetadataView::channels, &LogicalTransferMetadataView::sourceCores, &LogicalTransferMetadataView::targetCores};
|
||||
&LogicalTransferMetadataView::channels, &LogicalTransferMetadataView::sourceCores,
|
||||
&LogicalTransferMetadataView::targetCores};
|
||||
struct TransferGrids {
|
||||
std::array<StaticIntGrid, 3> values;
|
||||
StaticIntGrid &channels() { return values[0]; }
|
||||
@@ -41,7 +45,8 @@ template <typename Build> static FailureOr<TransferGrids> buildTransferGrids(Bui
|
||||
auto targetCores = build(transferMetadataMembers[2]);
|
||||
if (failed(channels) || failed(sourceCores) || failed(targetCores))
|
||||
return failure();
|
||||
return TransferGrids{{std::move(*channels), std::move(*sourceCores), std::move(*targetCores)}};
|
||||
return TransferGrids{{std::move(*channels), std::move(*sourceCores),
|
||||
std::move(*targetCores)}};
|
||||
}
|
||||
using GridGeometry = DeferredGridSliceGeometry;
|
||||
using StaticGeometryMember = SmallVector<StaticIntSequence> DeferredStaticSliceGeometry::*;
|
||||
@@ -82,6 +87,11 @@ static void appendMetadata(const ScheduledTransferSlice &slice, LogicalTransferM
|
||||
metadata.parentCounts.append(StaticIntSequence::uniform(family.requirement->exchange->externalTransferCount, count));
|
||||
metadata.sourceCores.append(family.sourceCores, familyIndex, count);
|
||||
metadata.targetCores.append(family.targetCores, familyIndex, count);
|
||||
if (family.hostRouted) {
|
||||
metadata.hostOffsets.append(family.hostOffsets, familyIndex, count);
|
||||
metadata.eventRegisters.append(
|
||||
family.eventRegisters, familyIndex, count);
|
||||
}
|
||||
metadata.targetLanes.append(StaticIntSequence::affine(targetLane, 1, count));
|
||||
if (family.requirement->producerLocalOffsets)
|
||||
metadata.localOffsets.append(*family.requirement->producerLocalOffsets, targetLane - requirementLanes.begin, count);
|
||||
@@ -172,6 +182,7 @@ static LogicalResult emitSendRun(const EmitSendRun &run, Value lane, unsigned la
|
||||
appendMetadata(slice, metadataByLane[sourceLane]);
|
||||
}
|
||||
LogicalTransferMetadataView logical = buildMetadataView(run.slices);
|
||||
ExternalTransferFamily &firstFamily = *run.slices.front().family;
|
||||
size_t actionCount = 0;
|
||||
for (const LogicalTransferMetadataView &laneMetadata : metadataByLane)
|
||||
actionCount = std::max(actionCount, laneMetadata.size());
|
||||
@@ -185,6 +196,20 @@ static LogicalResult emitSendRun(const EmitSendRun &run, Value lane, unsigned la
|
||||
FailureOr<StaticIntGrid> localOffsets = buildGrid(&LogicalTransferMetadataView::localOffsets, logical.localOffsets.valueAt(0));
|
||||
if (failed(transferGrids) || failed(localOffsets))
|
||||
return failure();
|
||||
std::optional<StaticIntGrid> hostOffsets;
|
||||
std::optional<StaticIntGrid> eventRegisters;
|
||||
if (firstFamily.hostRouted) {
|
||||
auto offsets = buildGrid(
|
||||
&LogicalTransferMetadataView::hostOffsets,
|
||||
logical.hostOffsets.valueAt(0));
|
||||
auto events = buildGrid(
|
||||
&LogicalTransferMetadataView::eventRegisters,
|
||||
logical.eventRegisters.valueAt(0));
|
||||
if (failed(offsets) || failed(events))
|
||||
return failure();
|
||||
hostOffsets = std::move(*offsets);
|
||||
eventRegisters = std::move(*events);
|
||||
}
|
||||
GridGeometry projectionGrids;
|
||||
for (auto [geometryIndex, sourceMember] : llvm::enumerate(metadataGeometryMembers)) {
|
||||
const auto &logicalValues = logical.*sourceMember;
|
||||
@@ -207,7 +232,6 @@ static LogicalResult emitSendRun(const EmitSendRun &run, Value lane, unsigned la
|
||||
const LogicalTransferMetadataView &source = metadataByLane[sourceLane];
|
||||
counts[sourceLane] = source.size();
|
||||
}
|
||||
ExternalTransferFamily &firstFamily = *run.slices.front().family;
|
||||
RequirementFamily &requirement = *firstFamily.requirement;
|
||||
Operation *anchor = requirement.exchange->deferred;
|
||||
Location loc = requirement.exchange->deferred.getLoc();
|
||||
@@ -217,10 +241,25 @@ static LogicalResult emitSendRun(const EmitSendRun &run, Value lane, unsigned la
|
||||
auto payload = materializeSendPayload(requirement, localOffset, projectionGrids[0].empty() ? nullptr : &projection, context, loc);
|
||||
if (failed(payload))
|
||||
return failure();
|
||||
auto send = SpatChannelSendOp::create(
|
||||
context.rewriter, loc, transferGrids->channels().emitLookup(action, runtimeLane, anchor, context.constants, context.rewriter, loc),
|
||||
transferGrids->sourceCores().emitLookup(action, runtimeLane, anchor, context.constants, context.rewriter, loc),
|
||||
transferGrids->targetCores().emitLookup(action, runtimeLane, anchor, context.constants, context.rewriter, loc), *payload);
|
||||
Value sourceCore = transferGrids->sourceCores().emitLookup(
|
||||
action, runtimeLane, anchor, context.constants, context.rewriter, loc);
|
||||
Value targetCore = transferGrids->targetCores().emitLookup(
|
||||
action, runtimeLane, anchor, context.constants, context.rewriter, loc);
|
||||
Operation *send;
|
||||
if (firstFamily.hostRouted)
|
||||
send = SpatHostStoreSyncOp::create(
|
||||
context.rewriter, loc, sourceCore, targetCore,
|
||||
hostOffsets->emitLookup(
|
||||
action, runtimeLane, anchor, context.constants, context.rewriter, loc),
|
||||
eventRegisters->emitLookup(
|
||||
action, runtimeLane, anchor, context.constants, context.rewriter, loc),
|
||||
*payload);
|
||||
else
|
||||
send = SpatChannelSendOp::create(
|
||||
context.rewriter, loc,
|
||||
transferGrids->channels().emitLookup(
|
||||
action, runtimeLane, anchor, context.constants, context.rewriter, loc),
|
||||
sourceCore, targetCore, *payload);
|
||||
setLogicalTransferMetadata(send, logical);
|
||||
return success();
|
||||
};
|
||||
@@ -255,14 +294,45 @@ static FailureOr<Value> emitReceiveValue(ArrayRef<ScheduledTransferSlice> slices
|
||||
};
|
||||
auto grids = buildTransferGrids([&](MetadataMember member) { return buildGrid(metadata.*member); });
|
||||
if (failed(grids)) return failure();
|
||||
std::optional<StaticIntGrid> hostOffsets;
|
||||
std::optional<StaticIntGrid> eventRegisters;
|
||||
if (slices.front().family->hostRouted) {
|
||||
auto offsets = buildGrid(metadata.hostOffsets);
|
||||
auto events = buildGrid(metadata.eventRegisters);
|
||||
if (failed(offsets) || failed(events))
|
||||
return failure();
|
||||
hostOffsets = std::move(*offsets);
|
||||
eventRegisters = std::move(*events);
|
||||
}
|
||||
Value position = lane ? lane : context.constants.getIndex(0);
|
||||
Value row = context.constants.getIndex(0);
|
||||
auto receive = SpatChannelReceiveOp::create(context.rewriter, anchor->getLoc(), requirement.publicationFragmentType,
|
||||
grids->channels().emitLookup(row, position, anchor, context.constants, context.rewriter, anchor->getLoc()),
|
||||
grids->sourceCores().emitLookup(row, position, anchor, context.constants, context.rewriter, anchor->getLoc()),
|
||||
grids->targetCores().emitLookup(row, position, anchor, context.constants, context.rewriter, anchor->getLoc()));
|
||||
Value sourceCore = grids->sourceCores().emitLookup(
|
||||
row, position, anchor, context.constants, context.rewriter, anchor->getLoc());
|
||||
Value targetCore = grids->targetCores().emitLookup(
|
||||
row, position, anchor, context.constants, context.rewriter, anchor->getLoc());
|
||||
Operation *receive;
|
||||
Value output;
|
||||
if (slices.front().family->hostRouted) {
|
||||
auto op = SpatHostWaitLoadOp::create(
|
||||
context.rewriter, anchor->getLoc(), requirement.publicationFragmentType,
|
||||
sourceCore, targetCore,
|
||||
hostOffsets->emitLookup(
|
||||
row, position, anchor, context.constants, context.rewriter, anchor->getLoc()),
|
||||
eventRegisters->emitLookup(
|
||||
row, position, anchor, context.constants, context.rewriter, anchor->getLoc()));
|
||||
receive = op;
|
||||
output = op.getOutput();
|
||||
} else {
|
||||
auto op = SpatChannelReceiveOp::create(
|
||||
context.rewriter, anchor->getLoc(), requirement.publicationFragmentType,
|
||||
grids->channels().emitLookup(
|
||||
row, position, anchor, context.constants, context.rewriter, anchor->getLoc()),
|
||||
sourceCore, targetCore);
|
||||
receive = op;
|
||||
output = op.getOutput();
|
||||
}
|
||||
setLogicalTransferMetadata(receive, metadata);
|
||||
return receive.getOutput();
|
||||
return output;
|
||||
}
|
||||
|
||||
static FailureOr<SmallVector<LogicalTransferMetadataView, 0>>
|
||||
@@ -315,6 +385,9 @@ static FailureOr<Value> emitReceiveAssembly(const EmitReceiveAssemblyRun &run, V
|
||||
SmallVector<int64_t> counts(laneCount);
|
||||
std::optional<TransferGrids> transferGrids;
|
||||
std::optional<StaticIntGrid> positions;
|
||||
std::optional<StaticIntGrid> hostOffsets;
|
||||
std::optional<StaticIntGrid> eventRegisters;
|
||||
bool hostRouted = run.slices.front().family->hostRouted;
|
||||
auto metadataByEntry = buildRectangularReceiveMetadata(run, laneCount);
|
||||
if (succeeded(metadataByEntry)) {
|
||||
auto buildRows = [&](auto member) {
|
||||
@@ -324,6 +397,16 @@ static FailureOr<Value> emitReceiveAssembly(const EmitReceiveAssemblyRun &run, V
|
||||
return StaticIntGrid::fromRows(rows);
|
||||
};
|
||||
auto grids = buildTransferGrids(buildRows);
|
||||
if (hostRouted) {
|
||||
auto offsets = buildRows(
|
||||
&LogicalTransferMetadataView::hostOffsets);
|
||||
auto events = buildRows(
|
||||
&LogicalTransferMetadataView::eventRegisters);
|
||||
if (failed(offsets) || failed(events))
|
||||
return failure();
|
||||
hostOffsets = std::move(*offsets);
|
||||
eventRegisters = std::move(*events);
|
||||
}
|
||||
SmallVector<StaticIntSequence> positionRows;
|
||||
for (unsigned position : run.positions)
|
||||
positionRows.push_back(StaticIntSequence::uniform(position, laneCount));
|
||||
@@ -368,6 +451,16 @@ static FailureOr<Value> emitReceiveAssembly(const EmitReceiveAssemblyRun &run, V
|
||||
return StaticIntGrid::fromColumns(actionCount, columns, defaultValue);
|
||||
};
|
||||
auto grids = buildTransferGrids(buildGrid);
|
||||
if (hostRouted) {
|
||||
auto offsets = buildGrid(
|
||||
&LogicalTransferMetadataView::hostOffsets);
|
||||
auto events = buildGrid(
|
||||
&LogicalTransferMetadataView::eventRegisters);
|
||||
if (failed(offsets) || failed(events))
|
||||
return failure();
|
||||
hostOffsets = std::move(*offsets);
|
||||
eventRegisters = std::move(*events);
|
||||
}
|
||||
SmallVector<StaticIntSequence> positionColumns;
|
||||
for (const StaticIntSequenceChain &values : positionsByLane)
|
||||
positionColumns.push_back(
|
||||
@@ -386,15 +479,34 @@ static FailureOr<Value> emitReceiveAssembly(const EmitReceiveAssemblyRun &run, V
|
||||
Value runtimeLane = lane ? lane : context.constants.getIndex(0);
|
||||
auto emitEntry = [&](Value entry, Value current) -> FailureOr<Value> {
|
||||
Type fragmentType = run.slices.front().family->requirement->publicationFragmentType;
|
||||
auto receive =
|
||||
SpatChannelReceiveOp::create(context.rewriter, loc, fragmentType,
|
||||
transferGrids->channels().emitLookup(entry, runtimeLane, anchor, context.constants, context.rewriter, loc),
|
||||
transferGrids->sourceCores().emitLookup(entry, runtimeLane, anchor, context.constants, context.rewriter, loc),
|
||||
transferGrids->targetCores().emitLookup(entry, runtimeLane, anchor, context.constants, context.rewriter, loc));
|
||||
Value sourceCore = transferGrids->sourceCores().emitLookup(
|
||||
entry, runtimeLane, anchor, context.constants, context.rewriter, loc);
|
||||
Value targetCore = transferGrids->targetCores().emitLookup(
|
||||
entry, runtimeLane, anchor, context.constants, context.rewriter, loc);
|
||||
Operation *receive;
|
||||
Value output;
|
||||
if (hostRouted) {
|
||||
auto op = SpatHostWaitLoadOp::create(
|
||||
context.rewriter, loc, fragmentType, sourceCore, targetCore,
|
||||
hostOffsets->emitLookup(
|
||||
entry, runtimeLane, anchor, context.constants, context.rewriter, loc),
|
||||
eventRegisters->emitLookup(
|
||||
entry, runtimeLane, anchor, context.constants, context.rewriter, loc));
|
||||
receive = op;
|
||||
output = op.getOutput();
|
||||
} else {
|
||||
auto op = SpatChannelReceiveOp::create(
|
||||
context.rewriter, loc, fragmentType,
|
||||
transferGrids->channels().emitLookup(
|
||||
entry, runtimeLane, anchor, context.constants, context.rewriter, loc),
|
||||
sourceCore, targetCore);
|
||||
receive = op;
|
||||
output = op.getOutput();
|
||||
}
|
||||
setLogicalTransferMetadata(receive, logical);
|
||||
Value position = positions->emitLookup(
|
||||
entry, runtimeLane, anchor, context.constants, context.rewriter, loc);
|
||||
return insert(receive.getOutput(), position, entry, runtimeLane, current);
|
||||
return insert(output, position, entry, runtimeLane, current);
|
||||
};
|
||||
if (actionCount == 1 && llvm::all_of(counts, [](int64_t count) { return count == 1; }))
|
||||
return emitEntry(context.constants.getIndex(0), initial);
|
||||
|
||||
+11
-12
@@ -28,6 +28,11 @@ static std::optional<Event> getPlannedHead(
|
||||
while (cursor.slice < plan.slices.size()) {
|
||||
const ScheduledTransferSlice &slice = plan.slices[cursor.slice];
|
||||
ExternalTransferFamily &family = *slice.family;
|
||||
if (family.hostRouted) {
|
||||
++cursor.slice;
|
||||
cursor.offset = 0;
|
||||
continue;
|
||||
}
|
||||
size_t begin = slice.familyOffset + cursor.offset;
|
||||
size_t length = slice.transferCount - cursor.offset;
|
||||
auto source = family.sourceStreams.find(stream, begin, length);
|
||||
@@ -243,6 +248,8 @@ LogicalResult verifyPlannedCommunicationDeadlockFree(
|
||||
DenseMap<ExternalTransferFamily *, unsigned> familyIndex;
|
||||
for (const ScheduledTransferSlice &slice : plan.slices) {
|
||||
ExternalTransferFamily *family = slice.family;
|
||||
if (family->hostRouted)
|
||||
continue;
|
||||
if (!familyIndex.try_emplace(family, familyIndex.size()).second)
|
||||
continue;
|
||||
size_t count = family->channelIds.size();
|
||||
@@ -258,18 +265,6 @@ LogicalResult verifyPlannedCommunicationDeadlockFree(
|
||||
familyChannels.emplace_back(
|
||||
first, first + static_cast<int64_t>(count));
|
||||
}
|
||||
llvm::sort(familyChannels);
|
||||
int64_t nextChannel = 0;
|
||||
for (auto [firstChannel, endChannel] : familyChannels) {
|
||||
if (firstChannel != nextChannel)
|
||||
return anchor->emitError(
|
||||
"planned communication channels are not exactly contiguous");
|
||||
nextChannel = endChannel;
|
||||
}
|
||||
if (static_cast<uint64_t>(nextChannel) != plan.logicalTransferCount)
|
||||
return anchor->emitError(
|
||||
"planned communication channel count is inconsistent");
|
||||
|
||||
for (const ScheduledTransferSlice &slice : plan.slices) {
|
||||
ExternalTransferFamily &family = *slice.family;
|
||||
for (size_t offset = 0; offset < slice.transferCount; ++offset) {
|
||||
@@ -296,6 +291,8 @@ LogicalResult verifyRealizedCommunicationDeadlockFree(
|
||||
DenseMap<ExternalTransferFamily *, unsigned> familyIndex;
|
||||
for (const ScheduledTransferSlice &slice : plan.slices) {
|
||||
ExternalTransferFamily *family = slice.family;
|
||||
if (family->hostRouted)
|
||||
continue;
|
||||
if (!familyIndex.try_emplace(family, familyIndex.size()).second)
|
||||
continue;
|
||||
for (size_t index = 0; index < family->channelIds.size(); ++index)
|
||||
@@ -305,6 +302,8 @@ LogicalResult verifyRealizedCommunicationDeadlockFree(
|
||||
DenseMap<int64_t, StaticIntSequenceChain> expected;
|
||||
for (const ScheduledTransferSlice &slice : plan.slices) {
|
||||
ExternalTransferFamily &family = *slice.family;
|
||||
if (family.hostRouted)
|
||||
continue;
|
||||
appendEventsByCore(expected, family.channelIds, family.sourceCores,
|
||||
slice.familyOffset, slice.transferCount, true);
|
||||
appendEventsByCore(expected, family.channelIds, family.targetCores,
|
||||
|
||||
+4
@@ -198,6 +198,7 @@ struct ScheduledInfo {
|
||||
llvm::SmallVector<mlir::Block*> blocks;
|
||||
llvm::SmallVector<mlir::Operation*> stepAnchors;
|
||||
llvm::SmallVector<int64_t> cores;
|
||||
llvm::SmallVector<unsigned> pipelineStages;
|
||||
unsigned stepCount = 0;
|
||||
llvm::SmallVector<ProducedValue*> produced;
|
||||
llvm::SmallVector<unsigned> streamIds;
|
||||
@@ -233,6 +234,9 @@ struct ExternalTransferFamily {
|
||||
StaticIntSequence sourceCores = StaticIntSequence::uniform(0, 1);
|
||||
StaticIntSequence targetCores = StaticIntSequence::uniform(0, 1);
|
||||
StaticIntSequence channelIds = StaticIntSequence::uniform(0, 1);
|
||||
StaticIntSequence hostOffsets = StaticIntSequence::uniform(0, 1);
|
||||
StaticIntSequence eventRegisters = StaticIntSequence::uniform(0, 1);
|
||||
bool hostRouted = false;
|
||||
};
|
||||
|
||||
struct DeferredExchangePlan {
|
||||
|
||||
+13
-2
@@ -209,15 +209,26 @@ static LogicalResult verifyDominance(func::FuncOp funcOp) {
|
||||
|
||||
LogicalResult realizeDeferredCommunication(func::FuncOp funcOp,
|
||||
const ScheduledComputeMaterializationResult& materialization,
|
||||
const SchedulingTarget& target) {
|
||||
const SchedulingTarget& target,
|
||||
size_t pipelineStages) {
|
||||
IRRewriter rewriter(funcOp.getContext());
|
||||
eraseUnusedIdentityDeferredCommunications(funcOp, rewriter);
|
||||
|
||||
auto transfers = buildDeferredTransferPlan(funcOp, materialization);
|
||||
auto transfers = buildDeferredTransferPlan(
|
||||
funcOp, materialization, pipelineStages, target.processorCount);
|
||||
if (failed(transfers))
|
||||
return funcOp.emitOpError("phase 2 failed to build symbolic transfer families");
|
||||
if (failed(placeLogicalProcessorsOnPhysicalCores(*transfers, target)))
|
||||
return failure();
|
||||
if (transfers->pipelineHostBufferBytes != 0) {
|
||||
auto bytes = pim::checkedCast<int64_t>(
|
||||
transfers->pipelineHostBufferBytes, funcOp,
|
||||
"pipeline host transfer storage");
|
||||
if (failed(bytes))
|
||||
return failure();
|
||||
funcOp->setAttr(kPipelineHostBufferBytesAttrName,
|
||||
rewriter.getI64IntegerAttr(*bytes));
|
||||
}
|
||||
auto schedule = scheduleDeferredCommunication(funcOp, *transfers);
|
||||
if (failed(schedule) || failed(verifyPlannedCommunicationDeadlockFree(funcOp, transfers->stepCounts, *schedule)))
|
||||
return funcOp.emitOpError("phase 2 failed to schedule symbolic communication");
|
||||
|
||||
+2
-1
@@ -9,6 +9,7 @@ struct SchedulingTarget;
|
||||
|
||||
mlir::LogicalResult realizeDeferredCommunication(mlir::func::FuncOp funcOp,
|
||||
const ScheduledComputeMaterializationResult& materialization,
|
||||
const SchedulingTarget& target);
|
||||
const SchedulingTarget& target,
|
||||
size_t pipelineStages = 1);
|
||||
|
||||
} // namespace onnx_mlir::spatial
|
||||
|
||||
+3
-2
@@ -11,7 +11,7 @@ using namespace mlir;
|
||||
namespace {
|
||||
|
||||
using TransferEmissionSignature =
|
||||
std::tuple<ScheduledInfo*, Value, Type, bool, bool, bool>;
|
||||
std::tuple<ScheduledInfo*, Value, Type, bool, bool, bool, bool>;
|
||||
|
||||
static TransferEmissionSignature getTransferEmissionSignature(
|
||||
const ExternalTransferFamily& family) {
|
||||
@@ -21,7 +21,8 @@ static TransferEmissionSignature getTransferEmissionSignature(
|
||||
family.requirement->publicationFragmentType,
|
||||
family.requirement->graphLanes.has_value(),
|
||||
family.requirement->producerProjection.has_value(),
|
||||
producer->scheduled->isBatch()};
|
||||
producer->scheduled->isBatch(),
|
||||
family.hostRouted};
|
||||
}
|
||||
|
||||
struct StreamThreshold {
|
||||
|
||||
+83
-14
@@ -5,6 +5,7 @@
|
||||
#include "DeferredProjectionAnalysis.hpp"
|
||||
#include "DeferredTransferPlanning.hpp"
|
||||
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
|
||||
#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
|
||||
|
||||
namespace onnx_mlir::spatial {
|
||||
using namespace mlir;
|
||||
@@ -28,7 +29,12 @@ static FailureOr<unsigned> getStepIndex(
|
||||
|
||||
static LogicalResult collectScheduledOperations(
|
||||
const ScheduledComputeMaterializationResult &materialization,
|
||||
DeferredTransferPlan &plan) {
|
||||
DeferredTransferPlan &plan,
|
||||
size_t pipelineStageCount,
|
||||
size_t processorCount) {
|
||||
if (pipelineStageCount == 0 || processorCount % pipelineStageCount != 0)
|
||||
return failure();
|
||||
size_t stageSize = processorCount / pipelineStageCount;
|
||||
unsigned nextStream = 0;
|
||||
for (const ScheduledMaterializationRecord &record :
|
||||
materialization.materializedSchedules) {
|
||||
@@ -46,8 +52,13 @@ static LogicalResult collectScheduledOperations(
|
||||
if (llvm::any_of(info.stepAnchors,
|
||||
[](Operation *anchor) { return !anchor; }))
|
||||
return op.emitOpError("phase 2 scheduled step anchor is missing");
|
||||
for (size_t core : record.cpus)
|
||||
for (size_t core : record.cpus) {
|
||||
if (core >= processorCount)
|
||||
return op.emitOpError("phase 2 scheduled core is outside the target");
|
||||
info.cores.push_back(core);
|
||||
if (pipelineStageCount > 1)
|
||||
info.pipelineStages.push_back(core / stageSize);
|
||||
}
|
||||
for (size_t lane = 0; lane < info.cores.size(); ++lane)
|
||||
info.streamIds.push_back(nextStream++);
|
||||
plan.scheduled.push_back(std::move(info));
|
||||
@@ -308,17 +319,22 @@ static LogicalResult buildRequirementFamilies(DeferredTransferPlan& plan,
|
||||
return success();
|
||||
}
|
||||
|
||||
static void buildAvailabilityFamilies(DeferredExchangePlan& exchange, uint64_t& nextChannel) {
|
||||
static LogicalResult buildAvailabilityFamilies(
|
||||
DeferredTransferPlan &plan,
|
||||
DeferredExchangePlan& exchange,
|
||||
uint64_t& nextChannel,
|
||||
DenseMap<int64_t, DenseMap<int64_t, unsigned>>& eventRegistersByTarget) {
|
||||
enum class Availability { Local, Direct, Host };
|
||||
for (RequirementFamily& requirement : exchange.requirements) {
|
||||
for (LaneInterval interval : requirement.targetLanes.intervals()) {
|
||||
unsigned runBegin = interval.begin;
|
||||
bool runLocal = false;
|
||||
Availability runAvailability = Availability::Local;
|
||||
bool haveRun = false;
|
||||
auto flush = [&](unsigned end) {
|
||||
auto flush = [&](unsigned end) -> LogicalResult {
|
||||
if (!haveRun || runBegin == end)
|
||||
return;
|
||||
return success();
|
||||
LaneSet lanes = LaneSet::range(runBegin, end);
|
||||
if (runLocal) {
|
||||
if (runAvailability == Availability::Local) {
|
||||
exchange.local.push_back({&requirement, lanes});
|
||||
}
|
||||
else {
|
||||
@@ -339,25 +355,72 @@ static void buildAvailabilityFamilies(DeferredExchangePlan& exchange, uint64_t&
|
||||
family.sourceCores = StaticIntSequence::uniform(requirement.producer->core, count);
|
||||
family.targetCores = StaticIntSequence::fromValues(targetCores);
|
||||
family.channelIds = StaticIntSequence::affine(nextChannel, 1, count);
|
||||
family.hostRouted = runAvailability == Availability::Host;
|
||||
if (family.hostRouted) {
|
||||
SmallVector<int64_t> eventRegisters;
|
||||
for (int64_t targetCore : targetCores) {
|
||||
auto ®isters = eventRegistersByTarget[targetCore];
|
||||
auto it = registers.try_emplace(
|
||||
requirement.producer->core, registers.size()).first;
|
||||
if (it->second >= kPimEventRegisterCount)
|
||||
return exchange.deferred.emitOpError(
|
||||
"pipeline host transfer requires more event registers than the target core provides");
|
||||
eventRegisters.push_back(it->second);
|
||||
}
|
||||
family.eventRegisters = StaticIntSequence::fromValues(
|
||||
eventRegisters);
|
||||
auto fragmentType = dyn_cast<ShapedType>(
|
||||
requirement.publicationFragmentType);
|
||||
auto fragmentBytes = fragmentType
|
||||
? pim::getCheckedShapedTypeSizeInBytes(
|
||||
fragmentType, exchange.deferred,
|
||||
"pipeline host transfer fragment")
|
||||
: FailureOr<uint64_t>(failure());
|
||||
if (failed(fragmentBytes))
|
||||
return failure();
|
||||
auto bytes = pim::checkedMul<size_t>(
|
||||
count, static_cast<size_t>(*fragmentBytes), exchange.deferred,
|
||||
"pipeline host transfer storage");
|
||||
if (failed(bytes))
|
||||
return failure();
|
||||
family.hostOffsets = StaticIntSequence::affine(
|
||||
plan.pipelineHostBufferBytes, *fragmentBytes, count);
|
||||
auto endOffset = pim::checkedAdd<size_t>(
|
||||
plan.pipelineHostBufferBytes, *bytes, exchange.deferred,
|
||||
"pipeline host transfer storage");
|
||||
if (failed(endOffset))
|
||||
return failure();
|
||||
plan.pipelineHostBufferBytes = *endOffset;
|
||||
}
|
||||
nextChannel += count;
|
||||
exchange.externalTransferCount += count;
|
||||
exchange.external.push_back(std::move(family));
|
||||
}
|
||||
return success();
|
||||
};
|
||||
for (unsigned lane = interval.begin; lane < interval.end; ++lane) {
|
||||
unsigned sourceStream = requirement.producer->scheduled->streamIds[requirement.producer->scheduledLane];
|
||||
bool local =
|
||||
sourceStream == exchange.target->streamIds[lane] && requirement.producer->step < exchange.consumerStep;
|
||||
if (haveRun && local != runLocal) {
|
||||
flush(lane);
|
||||
bool crossStage = !exchange.target->pipelineStages.empty()
|
||||
&& requirement.producer->scheduled->pipelineStages[
|
||||
requirement.producer->scheduledLane]
|
||||
!= exchange.target->pipelineStages[lane];
|
||||
Availability availability = local ? Availability::Local
|
||||
: crossStage ? Availability::Host : Availability::Direct;
|
||||
if (haveRun && availability != runAvailability) {
|
||||
if (failed(flush(lane)))
|
||||
return failure();
|
||||
runBegin = lane;
|
||||
}
|
||||
runLocal = local;
|
||||
runAvailability = availability;
|
||||
haveRun = true;
|
||||
}
|
||||
flush(interval.end);
|
||||
if (failed(flush(interval.end)))
|
||||
return failure();
|
||||
}
|
||||
}
|
||||
return success();
|
||||
}
|
||||
|
||||
static LogicalResult buildExchanges(func::FuncOp funcOp, DeferredTransferPlan& plan) {
|
||||
@@ -368,6 +431,7 @@ static LogicalResult buildExchanges(func::FuncOp funcOp, DeferredTransferPlan& p
|
||||
funcOp.walk([&](SpatDeferredCommunicationOp op) { deferredOps.push_back(op); });
|
||||
GraphBatchPublicationCache publicationCache;
|
||||
uint64_t nextChannel = 0;
|
||||
DenseMap<int64_t, DenseMap<int64_t, unsigned>> eventRegistersByTarget;
|
||||
for (SpatDeferredCommunicationOp deferred : deferredOps) {
|
||||
Operation* targetOp = deferred->getParentOfType<SpatScheduledCompute>();
|
||||
if (!targetOp)
|
||||
@@ -387,7 +451,9 @@ static LogicalResult buildExchanges(func::FuncOp funcOp, DeferredTransferPlan& p
|
||||
exchange->program = std::move(*program);
|
||||
if (failed(buildRequirementFamilies(plan, *exchange, publicationCache)))
|
||||
return failure();
|
||||
buildAvailabilityFamilies(*exchange, nextChannel);
|
||||
if (failed(buildAvailabilityFamilies(
|
||||
plan, *exchange, nextChannel, eventRegistersByTarget)))
|
||||
return failure();
|
||||
plan.exchanges.push_back(std::move(exchange));
|
||||
}
|
||||
return success();
|
||||
@@ -464,9 +530,12 @@ retargetBlueprint(DeferredTransferPlan& plan, SpatBlueprintOp blueprint, GraphBa
|
||||
|
||||
FailureOr<DeferredTransferPlan> buildDeferredTransferPlan(
|
||||
func::FuncOp funcOp,
|
||||
const ScheduledComputeMaterializationResult &materialization) {
|
||||
const ScheduledComputeMaterializationResult &materialization,
|
||||
size_t pipelineStages,
|
||||
size_t processorCount) {
|
||||
DeferredTransferPlan plan;
|
||||
if (failed(collectScheduledOperations(materialization, plan))
|
||||
if (failed(collectScheduledOperations(
|
||||
materialization, plan, pipelineStages, processorCount))
|
||||
|| failed(collectProducedValues(materialization, plan))
|
||||
|| failed(buildExchanges(funcOp, plan)))
|
||||
return failure();
|
||||
|
||||
+4
-1
@@ -13,11 +13,14 @@ struct DeferredTransferPlan {
|
||||
llvm::DenseMap<int64_t, llvm::SmallVector<ProducedValue*>> producedByGraph;
|
||||
llvm::SmallVector<std::unique_ptr<DeferredExchangePlan>> exchanges;
|
||||
llvm::SmallVector<unsigned> stepCounts;
|
||||
size_t pipelineHostBufferBytes = 0;
|
||||
};
|
||||
|
||||
mlir::FailureOr<DeferredTransferPlan>
|
||||
buildDeferredTransferPlan(mlir::func::FuncOp funcOp,
|
||||
const ScheduledComputeMaterializationResult &materialization);
|
||||
const ScheduledComputeMaterializationResult &materialization,
|
||||
size_t pipelineStages,
|
||||
size_t processorCount);
|
||||
|
||||
mlir::LogicalResult retargetDeferredPublications(mlir::func::FuncOp funcOp, DeferredTransferPlan& plan);
|
||||
|
||||
|
||||
+78
-7
@@ -3,12 +3,15 @@
|
||||
#include "DeferredCommunicationRealization.hpp"
|
||||
#include "ScheduledComputeReport.hpp"
|
||||
#include "ScheduledComputeVerification.hpp"
|
||||
#include "Scheduling/PipelineScheduling.hpp"
|
||||
#include "SpatialDataflowCsvExporter.hpp"
|
||||
#include "src/Accelerators/PIM/Common/PimCommon.hpp"
|
||||
#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Passes/Analyses/ONNXToSpatialVerifier.hpp"
|
||||
#include "src/Accelerators/PIM/Passes/PIMPasses.h"
|
||||
|
||||
#include <limits>
|
||||
|
||||
using namespace mlir;
|
||||
|
||||
namespace onnx_mlir {
|
||||
@@ -31,14 +34,47 @@ static FailureOr<func::FuncOp> requireEntry(ModuleOp moduleOp) {
|
||||
return *entry;
|
||||
}
|
||||
|
||||
static SchedulingTarget getPipelineSchedulingTarget(
|
||||
const SchedulingTarget& physicalTarget, size_t pipelineStages) {
|
||||
if (pipelineStages == 1)
|
||||
return physicalTarget;
|
||||
|
||||
SchedulingTarget schedulingTarget = physicalTarget;
|
||||
schedulingTarget.processorCount = physicalTarget.processorCount / pipelineStages;
|
||||
schedulingTarget.residentWeightCapacity = checkedMultiply(
|
||||
physicalTarget.residentWeightCapacity, pipelineStages);
|
||||
schedulingTarget.interProcessorLatencyNs.assign(
|
||||
schedulingTarget.processorCount * schedulingTarget.processorCount, 0);
|
||||
Cost latencySum = 0;
|
||||
size_t pairCount = 0;
|
||||
for (size_t source = 0; source < schedulingTarget.processorCount; ++source)
|
||||
for (size_t destination = 0;
|
||||
destination < schedulingTarget.processorCount; ++destination) {
|
||||
Cost latency = physicalTarget.getInterProcessorLatencyNs(
|
||||
source, destination);
|
||||
schedulingTarget.interProcessorLatencyNs[
|
||||
source * schedulingTarget.processorCount + destination] = latency;
|
||||
if (source != destination) {
|
||||
latencySum = checkedAdd(latencySum, latency);
|
||||
++pairCount;
|
||||
}
|
||||
}
|
||||
schedulingTarget.averageInterProcessorLatencyNs = pairCount == 0
|
||||
? 0
|
||||
: (latencySum + pairCount - 1) / pairCount;
|
||||
return schedulingTarget;
|
||||
}
|
||||
|
||||
struct ScheduleAndRealizeSpatialPass final
|
||||
: PassWrapper<ScheduleAndRealizeSpatialPass, OperationPass<ModuleOp>> {
|
||||
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ScheduleAndRealizeSpatialPass)
|
||||
|
||||
ScheduleAndRealizeSpatialPass() = default;
|
||||
ScheduleAndRealizeSpatialPass(const SchedulingTarget& target,
|
||||
SpatialDataflowExportStage exportStage)
|
||||
: target(target), exportStage(exportStage), hasTarget(true) {}
|
||||
SpatialDataflowExportStage exportStage,
|
||||
size_t pipelineStages)
|
||||
: target(target), exportStage(exportStage),
|
||||
pipelineStages(pipelineStages), hasTarget(true) {}
|
||||
|
||||
StringRef getArgument() const override { return "schedule-and-realize-spatial"; }
|
||||
StringRef getDescription() const override {
|
||||
@@ -52,6 +88,13 @@ struct ScheduleAndRealizeSpatialPass final
|
||||
signalPassFailure();
|
||||
return;
|
||||
}
|
||||
if (pipelineStages == 0 || target.processorCount % pipelineStages != 0
|
||||
|| target.residentWeightCapacity
|
||||
> std::numeric_limits<size_t>::max() / pipelineStages) {
|
||||
moduleOp.emitError("ScheduleAndRealizeSpatial requires valid pipeline stages and resource counts");
|
||||
signalPassFailure();
|
||||
return;
|
||||
}
|
||||
auto entry = requireEntry(moduleOp);
|
||||
if (failed(entry)) {
|
||||
signalPassFailure();
|
||||
@@ -59,8 +102,31 @@ struct ScheduleAndRealizeSpatialPass final
|
||||
}
|
||||
func::FuncOp entryFunc = *entry;
|
||||
|
||||
MergeSchedulingAnalysis analysis(entryFunc, target);
|
||||
MergeScheduleResult schedule = std::move(analysis.getResult());
|
||||
SchedulingTarget schedulingTarget = getPipelineSchedulingTarget(
|
||||
target, pipelineStages);
|
||||
ComputeGraph scheduledGraph;
|
||||
MergeScheduleResult schedule;
|
||||
for (;;) {
|
||||
MergeSchedulingAnalysis analysis(
|
||||
entryFunc, schedulingTarget,
|
||||
pipelineStages > 1 ? target.processorCount : 0);
|
||||
scheduledGraph = analysis.getGraph();
|
||||
schedule = std::move(analysis.getResult());
|
||||
std::string pipelineError;
|
||||
if (succeeded(applyPipelineScheduling(
|
||||
scheduledGraph, schedule, pipelineStages, target, pipelineError)))
|
||||
break;
|
||||
std::string splitError;
|
||||
if (pipelineStages == 1
|
||||
|| failed(splitPipelineWorkload(
|
||||
scheduledGraph, schedule, pipelineStages, target, splitError))) {
|
||||
if (!splitError.empty())
|
||||
pipelineError = splitError;
|
||||
moduleOp.emitError() << pipelineError;
|
||||
signalPassFailure();
|
||||
return;
|
||||
}
|
||||
}
|
||||
PatternRewriter rewriter(moduleOp.getContext());
|
||||
FailureOr<ScheduledComputeMaterializationResult> materialization =
|
||||
materializeScheduledCompute(entryFunc, schedule, rewriter);
|
||||
@@ -94,7 +160,8 @@ struct ScheduleAndRealizeSpatialPass final
|
||||
moduleOp, entryFunc, schedule, materializationResult.peftClassPlans,
|
||||
materializationResult.materializedSchedules);
|
||||
|
||||
if (failed(realizeDeferredCommunication(entryFunc, materializationResult, target))) {
|
||||
if (failed(realizeDeferredCommunication(
|
||||
entryFunc, materializationResult, target, pipelineStages))) {
|
||||
moduleOp.emitError("Spatial communication realization failed");
|
||||
signalPassFailure();
|
||||
return;
|
||||
@@ -126,6 +193,7 @@ struct ScheduleAndRealizeSpatialPass final
|
||||
private:
|
||||
SchedulingTarget target;
|
||||
SpatialDataflowExportStage exportStage = SpatialDataflowExportStage::None;
|
||||
size_t pipelineStages = 1;
|
||||
bool hasTarget = false;
|
||||
};
|
||||
|
||||
@@ -136,8 +204,11 @@ std::unique_ptr<Pass> createScheduleAndRealizeSpatialPass() {
|
||||
}
|
||||
|
||||
std::unique_ptr<Pass> createScheduleAndRealizeSpatialPass(
|
||||
const SchedulingTarget& target, SpatialDataflowExportStage exportStage) {
|
||||
return std::make_unique<ScheduleAndRealizeSpatialPass>(target, exportStage);
|
||||
const SchedulingTarget& target,
|
||||
SpatialDataflowExportStage exportStage,
|
||||
size_t pipelineStages) {
|
||||
return std::make_unique<ScheduleAndRealizeSpatialPass>(
|
||||
target, exportStage, pipelineStages);
|
||||
}
|
||||
|
||||
} // namespace spatial
|
||||
|
||||
+16
@@ -0,0 +1,16 @@
|
||||
#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
|
||||
+13
-4
@@ -772,6 +772,11 @@ std::vector<ComputeGraphEdge> aggregateEdges(llvm::ArrayRef<ComputeGraphEdge> ed
|
||||
|
||||
} // namespace
|
||||
|
||||
TransferCost getTransferCostFromBytes(Cost bytes,
|
||||
const SchedulingTarget& target) {
|
||||
return SchedulerCostModel {target}.getTransferCostFromBytes(bytes);
|
||||
}
|
||||
|
||||
uint64_t countComputeBodyInstructions(Region& body) {
|
||||
uint64_t numOperations = 0;
|
||||
body.walk([&](Operation* op) { numOperations = checkedAdd(numOperations, static_cast<uint64_t>(1)); });
|
||||
@@ -875,9 +880,13 @@ ResidentWeightSet getComputeInstanceResidentWeights(const ComputeInstance& insta
|
||||
return tiled;
|
||||
}
|
||||
|
||||
ComputeGraph buildComputeGraph(Operation* entryOp, const SchedulingTarget& target) {
|
||||
ComputeGraph buildComputeGraph(Operation* entryOp,
|
||||
const SchedulingTarget& target,
|
||||
size_t computePartitionCount) {
|
||||
ComputeGraph graph;
|
||||
SchedulerCostModel costModel {target};
|
||||
if (computePartitionCount == 0)
|
||||
computePartitionCount = target.processorCount;
|
||||
|
||||
for (Region& region : entryOp->getRegions()) {
|
||||
for (Block& block : region) {
|
||||
@@ -898,10 +907,10 @@ ComputeGraph buildComputeGraph(Operation* entryOp, const SchedulingTarget& targe
|
||||
if (isUsedAsWeightOnly(batch.getOperation()))
|
||||
continue;
|
||||
size_t chunkCount =
|
||||
getBatchChunkTargetCount(batch, target.processorCount);
|
||||
getBatchChunkTargetCount(batch, computePartitionCount);
|
||||
for (size_t chunkIndex = 0; chunkIndex < chunkCount; ++chunkIndex) {
|
||||
ComputeInstance instance = getBatchChunkForIndex(
|
||||
batch, chunkIndex, target.processorCount);
|
||||
batch, chunkIndex, computePartitionCount);
|
||||
size_t index = graph.nodes.size();
|
||||
graph.nodes.push_back({instance,
|
||||
getComputeInstanceCost(instance, target),
|
||||
@@ -920,7 +929,7 @@ ComputeGraph buildComputeGraph(Operation* entryOp, const SchedulingTarget& targe
|
||||
for (Value input : inputs) {
|
||||
for (const ProducerValueRef& producerRef :
|
||||
collectProducerValueRefs(input, node.instance,
|
||||
target.processorCount)) {
|
||||
computePartitionCount)) {
|
||||
auto producerIt = graph.instanceToIndex.find(producerRef.instance);
|
||||
if (producerIt == graph.instanceToIndex.end())
|
||||
continue;
|
||||
|
||||
+5
-1
@@ -61,9 +61,13 @@ struct ComputeGraph {
|
||||
llvm::DenseMap<ComputeInstance, size_t> instanceToIndex;
|
||||
};
|
||||
|
||||
ComputeGraph buildComputeGraph(mlir::Operation* entryOp, const SchedulingTarget& target);
|
||||
ComputeGraph buildComputeGraph(mlir::Operation* entryOp,
|
||||
const SchedulingTarget& target,
|
||||
size_t computePartitionCount = 0);
|
||||
bool verifyAcyclic(const ComputeGraph& graph);
|
||||
|
||||
TransferCost getTransferCostFromBytes(Cost bytes,
|
||||
const SchedulingTarget& target);
|
||||
uint64_t countComputeBodyInstructions(mlir::Region& body);
|
||||
uint64_t countComputeBodyOperationInstances(mlir::Region& body);
|
||||
Cost getComputeInstanceCost(const ComputeInstance& instance, const SchedulingTarget& target);
|
||||
|
||||
+4
-3
@@ -89,13 +89,14 @@ void verifySchedule(const ComputeGraph& graph,
|
||||
} // namespace
|
||||
|
||||
MergeSchedulingAnalysis::MergeSchedulingAnalysis(mlir::Operation* op,
|
||||
const SchedulingTarget& schedulingTarget)
|
||||
: entryOp(op), target(schedulingTarget) {
|
||||
const SchedulingTarget& schedulingTarget,
|
||||
size_t partitionCount)
|
||||
: entryOp(op), target(schedulingTarget), computePartitionCount(partitionCount) {
|
||||
result = run();
|
||||
}
|
||||
|
||||
MergeScheduleResult MergeSchedulingAnalysis::run() {
|
||||
ComputeGraph graph = buildComputeGraph(entryOp, target);
|
||||
graph = buildComputeGraph(entryOp, target, computePartitionCount);
|
||||
if (!verifyAcyclic(graph))
|
||||
llvm::report_fatal_error("merge scheduling: compute graph is cyclic");
|
||||
|
||||
|
||||
+7
-1
@@ -3,6 +3,7 @@
|
||||
#include "mlir/IR/Operation.h"
|
||||
|
||||
#include "MergeSchedule.hpp"
|
||||
#include "ComputeGraph.hpp"
|
||||
#include "SchedulingTarget.hpp"
|
||||
|
||||
namespace onnx_mlir {
|
||||
@@ -10,12 +11,17 @@ namespace spatial {
|
||||
|
||||
class MergeSchedulingAnalysis {
|
||||
public:
|
||||
MergeSchedulingAnalysis(mlir::Operation* op, const SchedulingTarget& target);
|
||||
MergeSchedulingAnalysis(mlir::Operation* op,
|
||||
const SchedulingTarget& target,
|
||||
size_t computePartitionCount = 0);
|
||||
MergeScheduleResult& getResult() { return result; }
|
||||
const ComputeGraph& getGraph() const { return graph; }
|
||||
|
||||
private:
|
||||
mlir::Operation* entryOp = nullptr;
|
||||
const SchedulingTarget& target;
|
||||
size_t computePartitionCount = 0;
|
||||
ComputeGraph graph;
|
||||
MergeScheduleResult result;
|
||||
|
||||
MergeScheduleResult run();
|
||||
|
||||
+1092
File diff suppressed because it is too large
Load Diff
+26
@@ -0,0 +1,26 @@
|
||||
#pragma once
|
||||
|
||||
#include "mlir/Support/LogicalResult.h"
|
||||
|
||||
#include <cstddef>
|
||||
#include <string>
|
||||
|
||||
#include "ComputeGraph.hpp"
|
||||
#include "MergeSchedule.hpp"
|
||||
#include "SchedulingTarget.hpp"
|
||||
|
||||
namespace onnx_mlir::spatial {
|
||||
|
||||
mlir::LogicalResult applyPipelineScheduling(const ComputeGraph& graph,
|
||||
MergeScheduleResult& schedule,
|
||||
size_t pipelineStages,
|
||||
const SchedulingTarget& physicalTarget,
|
||||
std::string& error);
|
||||
|
||||
mlir::LogicalResult splitPipelineWorkload(const ComputeGraph& graph,
|
||||
const MergeScheduleResult& schedule,
|
||||
size_t pipelineStages,
|
||||
const SchedulingTarget& physicalTarget,
|
||||
std::string& error);
|
||||
|
||||
} // namespace onnx_mlir::spatial
|
||||
@@ -550,7 +550,8 @@ def SpatChannelSendOp : SpatOp<"channel_send", []> {
|
||||
);
|
||||
|
||||
let assemblyFormat = [{
|
||||
$input `channel` $channelId `from` $sourceCoreId `to` $targetCoreId attr-dict `:` type($input)
|
||||
$input `channel` $channelId `from` $sourceCoreId `to` $targetCoreId
|
||||
attr-dict `:` type($input)
|
||||
}];
|
||||
}
|
||||
|
||||
@@ -568,7 +569,45 @@ def SpatChannelReceiveOp : SpatOp<"channel_receive", []> {
|
||||
);
|
||||
|
||||
let assemblyFormat = [{
|
||||
`channel` $channelId `from` $sourceCoreId `to` $targetCoreId attr-dict `:` type($output)
|
||||
`channel` $channelId `from` $sourceCoreId `to` $targetCoreId
|
||||
attr-dict `:` type($output)
|
||||
}];
|
||||
}
|
||||
|
||||
def SpatHostStoreSyncOp : SpatOp<"host_store_sync", []> {
|
||||
let summary = "Store a tensor to host memory and signal its consumer";
|
||||
|
||||
let arguments = (ins
|
||||
Index:$sourceCoreId,
|
||||
Index:$targetCoreId,
|
||||
Index:$hostOffset,
|
||||
Index:$eventRegister,
|
||||
SpatTensor:$input
|
||||
);
|
||||
|
||||
let assemblyFormat = [{
|
||||
$input `from` $sourceCoreId `to` $targetCoreId
|
||||
`host_offset` $hostOffset `event` $eventRegister attr-dict `:` type($input)
|
||||
}];
|
||||
}
|
||||
|
||||
def SpatHostWaitLoadOp : SpatOp<"host_wait_load", []> {
|
||||
let summary = "Wait for a producer and load its tensor from host memory";
|
||||
|
||||
let arguments = (ins
|
||||
Index:$sourceCoreId,
|
||||
Index:$targetCoreId,
|
||||
Index:$hostOffset,
|
||||
Index:$eventRegister
|
||||
);
|
||||
|
||||
let results = (outs
|
||||
SpatTensor:$output
|
||||
);
|
||||
|
||||
let assemblyFormat = [{
|
||||
`from` $sourceCoreId `to` $targetCoreId
|
||||
`host_offset` $hostOffset `event` $eventRegister attr-dict `:` type($output)
|
||||
}];
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,24 @@
|
||||
#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
|
||||
@@ -0,0 +1,37 @@
|
||||
#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
|
||||
@@ -0,0 +1,63 @@
|
||||
#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
|
||||
@@ -16,7 +16,8 @@ enum class SpatialDataflowExportStage;
|
||||
std::unique_ptr<mlir::Pass> createScheduleAndRealizeSpatialPass();
|
||||
std::unique_ptr<mlir::Pass> createScheduleAndRealizeSpatialPass(
|
||||
const SchedulingTarget& target,
|
||||
SpatialDataflowExportStage exportStage);
|
||||
SpatialDataflowExportStage exportStage,
|
||||
size_t pipelineStages = 1);
|
||||
}
|
||||
|
||||
std::unique_ptr<mlir::Pass> createONNXToSpatialPass();
|
||||
|
||||
Reference in New Issue
Block a user