nnicolosi/Raptor
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

better throughput scheduling with steady state cost model, now limited to spatial

Browse Source
This commit is contained in:
NiccoloN
2026-10-26 12:07:01 +01:00
parent fd493ae6e8
commit 4b53f6cf10
62 changed files with 4293 additions and 871 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/PIM/CMakeLists.txt
+30
View File
@@ -30,6 +30,31 @@ add_pim_unittest(PimMemoryLivenessPlannerTest
OMPimCompilerUtils
)
add_pim_unittest(CopyAccessPlanTest
CopyAccessPlanTest.cpp
LINK_LIBS PRIVATE OMPimCompilerUtils
)
add_test(NAME CopyAccessPlanChecksActive COMMAND CopyAccessPlanTest --check-failure)
set_tests_properties(CopyAccessPlanChecksActive PROPERTIES WILL_FAIL TRUE LABELS pim-unittest)
add_test(NAME RowMajorStrideOverflow
COMMAND "${PYTHON_EXECUTABLE}" "${CMAKE_CURRENT_SOURCE_DIR}/test_ShapeUtils.py"
--test "$<TARGET_FILE:CopyAccessPlanTest>")
set_tests_properties(RowMajorStrideOverflow PROPERTIES LABELS pim-unittest)
add_pim_unittest(ComputeInputMaterializationTest
ComputeInputMaterializationTest.cpp
LINK_LIBS PRIVATE OMPimCompilerUtils
)
add_pim_unittest(ProjectionMemoryCostTest
ProjectionMemoryCostTest.cpp
LINK_LIBS PRIVATE OMPimCompilerUtils
)
add_pim_unittest(PimAbstractExecutionTest
PimAbstractExecutionTest.cpp
LINK_LIBS PRIVATE OMPimCompilerUtils
)
add_pim_unittest(SpatialSchedulingTargetTest
SpatialSchedulingTargetTest.cpp
@@ -77,3 +102,8 @@ add_pim_unittest(PimPipelineSelectionTest
PimPipelineSelectionTest.cpp
LINK_LIBS PRIVATE OMPimCompilerUtils
)
add_test(NAME PimSpatialSelectionTest
COMMAND "${PYTHON_EXECUTABLE}" "${CMAKE_CURRENT_SOURCE_DIR}/test_PimSpatialSelection.py"
--compiler "$<TARGET_FILE:onnx-mlir>")
set_tests_properties(PimSpatialSelectionTest PROPERTIES LABELS pim-unittest)
test/PIM/ComputeInputMaterializationTest.cpp
+158
View File
@@ -0,0 +1,158 @@
#include "TestCheck.hpp"
#include <cmath>
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/Verifier.h"
#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
#include "src/Accelerators/PIM/Conversion/SpatialToPim/ComputeLikeRegionUtils.hpp"
#include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
using namespace mlir;
using namespace onnx_mlir;
using namespace onnx_mlir::spatial;
static void testZeroContent(OpBuilder &builder) {
auto type = RankedTensorType::get({4}, builder.getF32Type());
for (double value : {0.0, -0.0, 1.0, -1.0}) {
auto attribute = DenseElementsAttr::get(type, builder.getF32FloatAttr(value));
TEST_CHECK(isZeroBitPatternSplat(attribute) == (value == 0.0 && !std::signbit(value)));
}
auto integerType = RankedTensorType::get({4}, builder.getI32Type());
TEST_CHECK(isZeroBitPatternSplat(DenseElementsAttr::get(integerType, builder.getI32IntegerAttr(0))));
TEST_CHECK(!isZeroBitPatternSplat(DenseElementsAttr::get(integerType, builder.getI32IntegerAttr(1))));
TEST_CHECK(!isZeroBitPatternSplat(DenseElementsAttr::get(type, ArrayRef<float>{0, 1, 0, 0})));
TEST_CHECK(!isZeroBitPatternSplat(builder.getF32FloatAttr(0)));
TEST_CHECK(!isZeroBitPatternSplat({}));
}
static void testInputPlan(MLIRContext &context, bool batched) {
OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
OpBuilder builder(module->getBodyRegion());
auto loc = module->getLoc();
auto type = RankedTensorType::get({1, 4}, builder.getF32Type());
auto function = func::FuncOp::create(builder, loc, "test", builder.getFunctionType({type}, {}));
builder.setInsertionPointToStart(function.addEntryBlock());
Value constant = arith::ConstantOp::create(builder, loc,
DenseElementsAttr::get(type, builder.getF32FloatAttr(2)));
Value view = tensor::CastOp::create(builder, loc, type, constant);
SmallVector<Value> inputs{constant, function.getArgument(0), constant, view, function.getArgument(0)};
Operation *owner;
if (batched) {
auto batch = SpatScheduledComputeBatch::create(builder, loc, TypeRange{},
builder.getI32IntegerAttr(2), ValueRange{}, inputs);
batch->setAttr("coreIds", builder.getDenseI32ArrayAttr({0, 1}));
owner = batch;
} else {
auto single = SpatScheduledCompute::create(builder, loc, TypeRange{}, ValueRange{}, inputs);
single->setAttr("coreId", builder.getI32IntegerAttr(0));
owner = single;
}
auto *body = new Block;
owner->getRegion(0).push_back(body);
if (batched) body->addArgument(builder.getIndexType(), loc);
for (Value input : inputs) body->addArgument(input.getType(), loc);
builder.setInsertionPointToStart(body);
unsigned first = batched ? 1 : 0;
Value zero = arith::ConstantIndexOp::create(builder, loc, 0);
Value limit = arith::ConstantIndexOp::create(builder, loc, 1000000000000);
Value one = arith::ConstantIndexOp::create(builder, loc, 1);
auto loop = scf::ForOp::create(builder, loc, zero, limit, one);
builder.setInsertionPointToStart(loop.getBody());
SpatVAddOp::create(builder, loc, type, body->getArgument(first), body->getArgument(first + 1));
SpatVAddOp::create(builder, loc, type, body->getArgument(first + 3), body->getArgument(first));
builder.setInsertionPointAfter(loop);
auto lastUse = SpatVAddOp::create(builder, loc, type, body->getArgument(first + 2), body->getArgument(first));
auto lastFunctionUse = SpatVAddOp::create(builder, loc, type,
body->getArgument(first + 4), body->getArgument(first));
SpatYieldOp::create(builder, loc, ValueRange{});
builder.setInsertionPointAfter(owner);
func::ReturnOp::create(builder, loc);
TEST_CHECK(succeeded(verify(*module)));
// Body lowering can remove an input's final use before input materialization.
lastUse.erase();
lastFunctionUse.erase();
std::string before;
llvm::raw_string_ostream stream(before);
module->print(stream);
for (unsigned repeat = 0; repeat < 2; ++repeat) {
auto constantPlan = planComputeInputMaterialization(owner, 0);
auto runtimePlan = planComputeInputMaterialization(owner, 1);
auto unusedPlan = planComputeInputMaterialization(owner, 2);
auto viewPlan = planComputeInputMaterialization(owner, 3);
TEST_CHECK(succeeded(constantPlan) && succeeded(runtimePlan) && succeeded(unusedPlan));
TEST_CHECK(*constantPlan == (batched ? ComputeInputMaterialization::CopyAtEntry
: ComputeInputMaterialization::ForwardConstant));
TEST_CHECK(*runtimePlan == ComputeInputMaterialization::ForwardFunctionInput);
auto unusedFunctionPlan = planComputeInputMaterialization(owner, 4);
TEST_CHECK(succeeded(unusedFunctionPlan)
&& *unusedFunctionPlan == (batched ? ComputeInputMaterialization::ForwardFunctionInput
: ComputeInputMaterialization::Unused));
TEST_CHECK(*unusedPlan == (batched ? ComputeInputMaterialization::CopyAtEntry
: ComputeInputMaterialization::Unused));
TEST_CHECK(succeeded(viewPlan) && *viewPlan == ComputeInputMaterialization::CopyAtEntry);
}
TEST_CHECK(failed(planComputeInputMaterialization(owner, 5)));
TEST_CHECK(failed(planComputeInputMaterialization(function, 0)));
std::string after;
llvm::raw_string_ostream afterStream(after);
module->print(afterStream);
TEST_CHECK(before == after);
}
static void testOutputReuse(MLIRContext &context) {
OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
IRRewriter builder(&context);
builder.setInsertionPointToStart(module->getBody());
auto loc = module->getLoc();
auto type = RankedTensorType::get({1, 4}, builder.getF32Type());
auto function = func::FuncOp::create(builder, loc, "reuse", builder.getFunctionType({}, {}));
builder.setInsertionPointToStart(function.addEntryBlock());
Value input = tensor::EmptyOp::create(builder, loc, type.getShape(), type.getElementType());
Value constant = arith::ConstantOp::create(builder, loc,
DenseElementsAttr::get(type, builder.getF32FloatAttr(2)));
auto first = SpatVAddOp::create(builder, loc, type, input, constant);
TEST_CHECK(getReusableOutputTensor(first) == input);
auto second = SpatVAddOp::create(builder, loc, type, input, constant);
TEST_CHECK(!getReusableOutputTensor(first));
TEST_CHECK(getReusableOutputTensor(second) == input);
second->moveBefore(first);
TEST_CHECK(getReusableOutputTensor(first) == input);
TEST_CHECK(!getReusableOutputTensor(second));
second->moveAfter(first);
TEST_CHECK(!getReusableOutputTensor(first));
TEST_CHECK(getReusableOutputTensor(second) == input);
Value view = tensor::CastOp::create(builder, loc, type, input);
auto fromView = SpatVAddOp::create(builder, loc, type, view, constant);
TEST_CHECK(!getReusableOutputTensor(fromView));
TEST_CHECK(getBestOutputTensorFromOperandsOrAllocate(builder, fromView)
.getDefiningOp<tensor::EmptyOp>());
builder.setInsertionPointAfter(fromView);
Value zero = arith::ConstantIndexOp::create(builder, loc, 0);
Value limit = arith::ConstantIndexOp::create(builder, loc, 10);
Value one = arith::ConstantIndexOp::create(builder, loc, 1);
auto loop = scf::ForOp::create(builder, loc, zero, limit, one, ValueRange{input});
builder.setInsertionPointToStart(loop.getBody());
auto carried = SpatVAddOp::create(builder, loc, type, loop.getRegionIterArg(0), constant);
scf::YieldOp::create(builder, loc, ValueRange{carried});
TEST_CHECK(getReusableOutputTensor(carried) == loop.getRegionIterArg(0));
TEST_CHECK(getBestOutputTensorFromOperandsOrAllocate(builder, carried) == loop.getRegionIterArg(0));
builder.setInsertionPointAfter(loop);
func::ReturnOp::create(builder, loc);
TEST_CHECK(succeeded(verify(*module)));
}
int main() {
MLIRContext context(MLIRContext::Threading::DISABLED);
context.loadDialect<arith::ArithDialect, func::FuncDialect, scf::SCFDialect,
tensor::TensorDialect, SpatialDialect>();
OpBuilder builder(&context);
testZeroContent(builder);
testInputPlan(context, false);
testInputPlan(context, true);
testOutputReuse(context);
}
test/PIM/CopyAccessPlanTest.cpp
+96
View File
@@ -0,0 +1,96 @@
#include "src/Accelerators/PIM/Common/IR/CopyAccessPlan.hpp"
#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/SchedulingTarget.hpp"
#include "TestCheck.hpp"
#include <limits>
#include <random>
using namespace onnx_mlir;
int main(int argc, char **argv) {
if (argc > 1 && std::string(argv[1]) == "--stride-overflow") {
computeRowMajorStrides({2, std::numeric_limits<int64_t>::max(), 2});
return 0;
}
if (argc > 1) TEST_CHECK(false);
auto dense = CopyAccessPlan::build({3, 4}, {4, 1}, {4, 1});
TEST_CHECK(mlir::succeeded(dense) && dense->getRunCount() == 1 && dense->getRunElements() == 12);
auto strided = CopyAccessPlan::build({3, 4}, {8, 1}, {4, 1});
TEST_CHECK(mlir::succeeded(strided) && strided->getRunCount() == 3 && strided->getRunElements() == 4);
TEST_CHECK(*strided->getRunOffsets(2) == std::make_pair(int64_t{16}, int64_t{8}));
TEST_CHECK(mlir::failed(strided->getRunOffsets(-1)) && mlir::failed(strided->getRunOffsets(3)));
auto large = CopyAccessPlan::build({1000000000, 16}, {32, 1}, {16, 1});
TEST_CHECK(mlir::succeeded(large) && large->getOuterShape().size() == 1);
TEST_CHECK(large->getRunCount() == 1000000000 && large->getRunElements() == 16);
TEST_CHECK(*large->getRunOffsets(999999999) == std::make_pair(int64_t{31999999968}, int64_t{15999999984}));
auto scalar = CopyAccessPlan::build({}, {}, {});
TEST_CHECK(mlir::succeeded(scalar) && scalar->getRunCount() == 1 && scalar->getRunElements() == 1);
// Observed input planes and channel-wise publication slices, without PIM IR.
auto input = CopyAccessPlan::build({3, 1, 224}, {50176, 224, 1}, {224, 224, 1});
TEST_CHECK(mlir::succeeded(input) && input->getRunCount() == 3 && input->getRunElements() == 224);
TEST_CHECK(*input->getRunOffsets(1) == std::make_pair(int64_t{50176}, int64_t{224}));
for (int64_t channels : {32, 96}) {
auto publication = CopyAccessPlan::build({1, channels, 1, 80},
{channels * 6400, 6400, 80, 1}, {channels * 80, 80, 80, 1});
TEST_CHECK(mlir::succeeded(publication) && publication->getRunCount() == channels);
TEST_CHECK(publication->getRunElements() == 80);
TEST_CHECK(*publication->getRunOffsets(1) == std::make_pair(int64_t{6400}, int64_t{80}));
}
TEST_CHECK(mlir::failed(CopyAccessPlan::build({2}, {}, {1})));
TEST_CHECK(mlir::failed(CopyAccessPlan::build({0}, {1}, {1})));
TEST_CHECK(mlir::failed(CopyAccessPlan::build({-1}, {1}, {1})));
TEST_CHECK(mlir::failed(CopyAccessPlan::build({2}, {-1}, {1})));
int64_t limit = std::numeric_limits<int64_t>::max();
TEST_CHECK(mlir::failed(computeRowMajorStridesChecked({2, limit, 2})));
auto limitStrides = computeRowMajorStridesChecked({2, limit});
TEST_CHECK(mlir::succeeded(limitStrides) && (*limitStrides)[0] == limit);
TEST_CHECK(mlir::failed(CopyAccessPlan::build({limit, 2}, {2, 1}, {2, 1})));
TEST_CHECK(mlir::failed(CopyAccessPlan::build({limit, 2}, {3, 1}, {3, 1})));
TEST_CHECK(mlir::failed(CopyAccessPlan::build({3}, {limit}, {1})));
TEST_CHECK(mlir::failed(CopyAccessPlan::build({3}, {1}, {limit})));
// A run start can fit even when its final element does not.
TEST_CHECK(mlir::failed(CopyAccessPlan::build({2, 2}, {limit, 1}, {2, 1})));
TEST_CHECK(mlir::failed(CopyAccessPlan::build({2, 2}, {2, 1}, {limit, 1})));
auto boundary = CopyAccessPlan::build({2}, {limit}, {1});
TEST_CHECK(mlir::succeeded(boundary) && boundary->getRunOffsets(1)->first == limit);
auto unit = CopyAccessPlan::build({1, 2}, {limit, 1}, {-1, 1});
TEST_CHECK(mlir::succeeded(unit) && unit->getRunElements() == 2);
// Independent element-by-element reference checks ordered offsets, not totals.
std::mt19937 random(123);
for (unsigned trial = 0; trial < 2000; ++trial) {
size_t rank = random() % 4 + 1;
llvm::SmallVector<int64_t> shape(rank), source(rank), target(rank);
for (auto &size : shape) size = random() % 4 + 1;
auto rowMajor = computeRowMajorStrides(shape);
for (size_t dim = 0; dim < rank; ++dim) {
source[dim] = random() % 2 ? rowMajor[dim] : random() % 20;
target[dim] = random() % 2 ? rowMajor[dim] : random() % 20;
}
auto plan = CopyAccessPlan::build(shape, source, target);
TEST_CHECK(mlir::succeeded(plan));
int64_t linear = 0;
for (int64_t run = 0; run < plan->getRunCount(); ++run) {
auto offsets = plan->getRunOffsets(run);
TEST_CHECK(mlir::succeeded(offsets));
for (int64_t element = 0; element < plan->getRunElements(); ++element, ++linear) {
auto indices = delinearizeIndex(linear, shape, rowMajor);
TEST_CHECK(offsets->first + element == linearizeIndex(indices, source));
TEST_CHECK(offsets->second + element == linearizeIndex(indices, target));
}
}
TEST_CHECK(linear == getNumElements(shape));
}
// The cost caller and lowering use the same geometry; target timing stays separate.
for (int64_t width : {8, 16}) {
spatial::SchedulingTarget target;
target.sharedMemoryWidthBytes = width;
target.sharedMemoryReadLatencyNs = width == 8 ? 3 : 8;
auto access = CopyAccessPlan::build({3, 5}, {8, 1}, {5, 1});
TEST_CHECK(mlir::succeeded(access));
auto demand = target.getSharedMemoryDemand(access->getRunElements() * 4, access->getRunCount(), false);
TEST_CHECK(demand.readBytes == 60 && demand.readRequests == 3);
TEST_CHECK(demand.serviceNs == 3 * ((20 + width - 1) / width) * target.sharedMemoryReadLatencyNs);
}
}
test/PIM/PimAbstractExecutionTest.cpp
+864
View File
@@ -0,0 +1,864 @@
#include "TestCheck.hpp"
#include <limits>
#include "src/Accelerators/PIM/Compiler/PimPeriodicAnalysis.hpp"
#include "src/Accelerators/PIM/Compiler/PimInstructionTiming.hpp"
#include "src/Accelerators/PIM/Compiler/PimSpatialCost.hpp"
#include "src/Accelerators/PIM/Compiler/PimBinaryFormat.hpp"
#include "src/Accelerators/PIM/Compiler/PimDependencyTokens.hpp"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/IR/Verifier.h"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
using namespace onnx_mlir;
using namespace mlir;
using E = PimExecutionEvent;
using O = pim_binary::Opcode;
static PimInstructionTiming hardware(int64_t period) {
PimInstructionTiming target;
target.processorCount = 2;
target.processorPeriodNs = period;
target.localMemoryPeriodNs = 3;
target.peerPermissionThenSourceRead = true;
target.signalIssueNs = period;
target.controlMessageBytes = 1;
target.sharedMemoryWidthBytes = 64;
target.sharedMemoryReadLatencyNs = 80;
target.sharedMemoryWriteLatencyNs = 120;
target.sharedMemoryRequestLatencyNs = {2, 3};
target.sharedMemoryResponseLatencyNs = {2, 3};
target.interProcessorLatencyNs = {0, 2, 2, 0};
target.synchronizationRegisterCount = 4;
return target;
}
static std::string binary(const std::vector<pim_binary::InstructionRecord> &records) {
std::string result(12, '\0');
result.replace(0, 4, "PIMB");
llvm::support::endian::write32le(result.data() + 4, 1);
llvm::support::endian::write32le(result.data() + 8, records.size());
for (const auto &record : records) {
auto encoded = pim_binary::encodeInstructionRecord(record);
result.append(encoded.data(), encoded.size());
}
return result;
}
static void testTypedResult(PimInstructionTiming target) {
PimPeriodicOptions options;
options.verifyPeriodicState = true;
PimExecutionProgram program{{{E::Work, 29}, {E::Load, 65}}};
auto estimate = evaluatePimExecution(program, target, options);
TEST_CHECK(estimate.cost.available() && estimate.cost.convergenceProved() == true);
TEST_CHECK(estimate.report.getString("status") == "estimated");
TEST_CHECK(estimate.cost.loweringSupport == spatial::StaticLoweringSupport::Unknown);
TEST_CHECK(!estimate.cost.eligible() && !estimate.cost.selectionMargin);
auto *queue = estimate.report.getObject("memory_controller");
for (size_t policy = 0; policy < 2; ++policy) {
auto *state = queue->getObject(policy ? "alternative_periodic_state" : "periodic_state");
const auto &proof = estimate.cost.recurrence[policy];
TEST_CHECK(proof.proved() && state->getString("status") == "proved");
TEST_CHECK(proof.intervalNs == state->getNumber("interval_ns"));
TEST_CHECK(proof.periodGenerations == state->getInteger("period_generations"));
}
auto primary = estimate.cost.primary, alternative = estimate.cost.alternative;
estimate.report = llvm::json::Object{{"status", "unavailable"}, {"convergence_proved", false}};
TEST_CHECK(estimate.cost.primary == primary && estimate.cost.alternative == alternative);
TEST_CHECK(estimate.cost.convergenceProved() == true);
options.verifyPeriodicState = false;
estimate = evaluatePimExecution(program, target, options);
TEST_CHECK(estimate.cost.available() && !estimate.cost.convergenceProved().has_value());
TEST_CHECK(estimate.report.getString("status") == "estimated");
TEST_CHECK(!estimate.report.getObject("memory_controller")->getObject("periodic_state"));
options.intervalWindowEventBudget = 0;
for (bool detailed : {false, true}) {
options.detailed = detailed;
estimate = evaluatePimExecution(program, target, options);
TEST_CHECK(estimate.queue && !estimate.cost.available());
TEST_CHECK(!estimate.cost.convergenceProved().has_value());
TEST_CHECK(estimate.report.getString("status") == "unavailable");
TEST_CHECK(estimate.report.getString("reason") == estimate.cost.reason);
TEST_CHECK(estimate.report.getObject("memory_controller")->getString("status") == "unavailable");
TEST_CHECK(!estimate.report.getObject("memory_controller")->getObject("periodic_state"));
}
options.intervalWindowEventBudget.reset();
options.detailed = false;
options.verifyPeriodicState = true;
estimate = evaluatePimExecution({{{E::Wait, 0, 0, 0, 1}, {E::Load, 65}}}, target, options);
TEST_CHECK(!estimate.cost.available() && estimate.cost.convergenceProved() == false);
TEST_CHECK(!estimate.cost.reason.empty() && !estimate.cost.eligible());
target.sharedMemoryWidthBytes = 0;
estimate = evaluatePimExecution(program, target, options);
TEST_CHECK(!estimate.cost.available() && !estimate.cost.reason.empty());
TEST_CHECK(estimate.report.getString("status") == "unavailable");
TEST_CHECK(estimate.report.getString("reason") == estimate.cost.reason);
}
static void testLateRecurrenceAuthority(PimInstructionTiming target) {
PimPeriodicOptions options;
options.verifyPeriodicState = true;
options.intervalWindowEventBudget = 1;
PimExecutionProgram program{{{E::Work, 29}, {E::Load, 65}}};
auto late = evaluatePimExecution(program, target, options);
TEST_CHECK(late.queue && late.cost.available() && late.cost.convergenceProved() == true);
for (const auto &policy : late.queue->policies) {
TEST_CHECK(policy.legacy.termination == PimQueueTermination::EventBudget);
TEST_CHECK(!policy.legacy.intervalNs && policy.recurrence.proved());
TEST_CHECK(policy.proofProcessedEvents > policy.legacy.progress.processedEvents);
TEST_CHECK(policy.intervalNs() == policy.recurrence.intervalNs);
TEST_CHECK(policy.termination == PimQueueTermination::RecurrenceProved);
}
auto first = serializePimQueueEvaluation(*late.queue);
auto second = serializePimQueueEvaluation(*late.queue);
TEST_CHECK(first.getString("status") == "estimated" && !first.getString("reason"));
TEST_CHECK(first.getString("alternative_tie_status") == "estimated");
TEST_CHECK(first.getNumber("interval_ns") == late.cost.primary);
TEST_CHECK(first.getNumber("alternative_tie_interval_ns") == late.cost.alternative);
TEST_CHECK(first.getObject("legacy_interval_window")->getString("status") == "unavailable");
TEST_CHECK(first.getObject("legacy_interval_window")->getString("termination") == "event-budget");
TEST_CHECK(llvm::json::Value(std::move(first)) == llvm::json::Value(std::move(second)));
auto blocked = evaluatePimExecution({{{E::Wait, 0, 0, 0, 1}, {E::Signal, 0, 1}},
{{E::Wait, 0, 0, 0, 1}, {E::Signal, 0, 0}}}, target, options);
TEST_CHECK(blocked.queue && blocked.cost.convergenceProved() == false);
TEST_CHECK(blocked.queue->policies[1].termination == PimQueueTermination::NoReadyEvents);
for (size_t failedPolicy = 0; failedPolicy < 2; ++failedPolicy) {
auto asymmetric = *late.queue;
asymmetric.policies[failedPolicy] = blocked.queue->policies[failedPolicy];
auto cost = asymmetric.cost();
cost.loweringSupport = spatial::StaticLoweringSupport::Supported;
cost.selectionMargin = 0.021;
TEST_CHECK(cost.primary.has_value() == (failedPolicy != 0));
TEST_CHECK(cost.alternative.has_value() == (failedPolicy != 1));
TEST_CHECK(cost.convergenceProved() == false && !cost.eligible());
TEST_CHECK(cost.reason == blocked.queue->policies[failedPolicy].recurrence.reason);
auto report = serializePimQueueEvaluation(asymmetric);
TEST_CHECK(report.getString("status") == (failedPolicy == 0 ? "unavailable" : "estimated"));
TEST_CHECK(report.getString("alternative_tie_status") == (failedPolicy == 1 ? "unavailable" : "estimated"));
TEST_CHECK(report.getObject(failedPolicy == 0 ? "periodic_state" : "alternative_periodic_state")
->getString("status") == "unproved");
}
options.intervalWindowEventBudget.reset();
auto normal = evaluatePimExecution(program, target, options);
options.workBreakdown = true;
auto detailed = evaluatePimExecution(program, target, options);
TEST_CHECK(late.cost.primary == normal.cost.primary && normal.cost.primary == detailed.cost.primary);
TEST_CHECK(late.cost.alternative == normal.cost.alternative && normal.cost.alternative == detailed.cost.alternative);
TEST_CHECK(normal.queue->policies[0].progress.processedEvents
<= detailed.queue->policies[0].progress.processedEvents);
}
static void testSharedEngine(PimInstructionTiming target) {
auto record = [](O op, int32_t peer = 0, int32_t reg = 0, int32_t count = 0, int32_t bytes = 0) {
pim_binary::InstructionRecord r;
r.opcode = op; r.r2OrImm = peer; r.generic1 = reg; r.generic2 = count; r.generic3 = bytes;
return r;
};
auto vectorWork = getPimVectorWorkNs(target, 65, 1, 1, true);
PimInstructionProgram instructions;
instructions.instructions = {
binary({record(O::setbw, 0, 8, 8), record(O::ld, 0, 0, 0, 65),
record(O::vvadd, 0, 0, 0, 65), record(O::send, 1, 0, 0, 65), record(O::sync, 1, 0)}),
binary({record(O::recv, 0, 0, 0, 65), record(O::wait, 0, 0, 1), record(O::st, 0, 0, 0, 33)})};
instructions.arrayGroups.resize(2);
PimExecutionProgram abstract(2);
abstract[0] = {{E::Work, target.processorPeriodNs}, {E::Load, 65}, {E::Work, vectorWork},
{E::Send, 65, 1}, {E::Signal, 0, 1, 0}};
abstract[1] = {{E::Receive, 65, 0}, {E::Wait, 0, 0, 0, 1}, {E::Store, 33}};
PimPeriodicOptions details;
details.workBreakdown = true;
PimExecutionProgram trace;
details.executionTrace = &trace;
auto reference = analyzePimInstructions(instructions, target, 2, details);
TEST_CHECK(reference.getString("status") == "estimated");
TEST_CHECK(!reference.getObject("memory_controller")->getObject("periodic_state"));
details.executionTrace = nullptr;
auto roundTrip = analyzePimExecution(trace, target, details);
for (const char *field : {"interval_ns", "alternative_tie_interval_ns"})
TEST_CHECK(roundTrip.getObject("memory_controller")->getNumber(field)
== reference.getObject("memory_controller")->getNumber(field));
for (size_t resource = 0; resource < 7; ++resource)
TEST_CHECK((*roundTrip.getArray("resource_service_ns"))[resource].getAsInteger()
== (*reference.getArray("resource_service_ns"))[resource].getAsInteger());
details.verifyPeriodicState = true;
auto proof = analyzePimExecution(trace, target, details);
TEST_CHECK(proof.getObject("memory_controller")->getString("status")
== reference.getObject("memory_controller")->getString("status"));
TEST_CHECK(proof.getObject("memory_controller")->getObject("legacy_interval_window")->getNumber("interval_ns")
== reference.getObject("memory_controller")->getNumber("interval_ns"));
TEST_CHECK(proof.getObject("memory_controller")->getNumber("interval_ns")
== proof.getObject("memory_controller")->getObject("periodic_state")->getNumber("interval_ns"));
TEST_CHECK(proof.getObject("memory_controller")->getObject("periodic_state"));
details.verifyPeriodicState = false;
auto actual = analyzePimExecution(abstract, target, details);
for (size_t core = 0; core < 2; ++core)
TEST_CHECK((*actual.getArray("request_order"))[core].getAsObject()->getString("identity")
== (*reference.getArray("request_order"))[core].getAsObject()->getString("identity"));
for (size_t core = 0; core < 2; ++core)
TEST_CHECK((*actual.getArray("request_order"))[core].getAsObject()->getString("blocking_identity")
== (*reference.getArray("request_order"))[core].getAsObject()->getString("blocking_identity"));
for (const char *field : {"global_read_bytes", "global_write_bytes", "global_read_requests",
"global_write_requests", "global_memory_service_ns"})
TEST_CHECK(actual.getInteger(field) == reference.getInteger(field));
auto *a = actual.getObject("memory_controller"), *b = reference.getObject("memory_controller");
TEST_CHECK(a && b && a->getString("status") == "estimated");
TEST_CHECK(a->getNumber("interval_ns") == b->getNumber("interval_ns"));
TEST_CHECK(a->getNumber("alternative_tie_interval_ns") == b->getNumber("alternative_tie_interval_ns"));
abstract[1].clear();
abstract[0] = {{E::Load, 1}, {E::Load, 65}};
auto ordered = analyzePimExecution(abstract, target, details);
abstract[0] = {{E::Load, 65}, {E::Load, 1}};
auto reversed = analyzePimExecution(abstract, target, details);
TEST_CHECK(ordered.getInteger("global_read_bytes") == reversed.getInteger("global_read_bytes"));
TEST_CHECK((*ordered.getArray("request_order"))[0].getAsObject()->getString("identity")
!= (*reversed.getArray("request_order"))[0].getAsObject()->getString("identity"));
// Preserve repeated request order, not a single larger request.
abstract[0] = {{E::Repeat, 2, 0, 0, 0, {{E::Load, 1}, {E::Work, 4}}}};
abstract[1].clear();
actual = analyzePimExecution(abstract, target);
TEST_CHECK(actual.getInteger("global_read_requests") == 2);
TEST_CHECK(actual.getInteger("global_memory_service_ns") == 160);
abstract[0] = {{E::Repeat, 1000000000000ull, 0, 0, 0, {{E::Load, 1}}}};
actual = analyzePimExecution(abstract, target);
TEST_CHECK(actual.getString("status") == "unavailable");
abstract[0] = {{E::Repeat, 1000000000000ull, 0, 0, 0, {{E::Work, 4}}}};
actual = analyzePimExecution(abstract, target);
TEST_CHECK(actual.getObject("memory_controller")->getNumber("interval_ns") == 4000000000000.0);
abstract[0] = {{E::Signal, 0, int64_t(1) << 32}};
TEST_CHECK(analyzePimExecution(abstract, target).getString("status") == "unavailable");
TEST_CHECK(getPimMatrixOutputBytes(target, 2, pim_binary::kMatrixWeightPrecisionBits, 3, 1)
== target.matrixColumns * 2 / pim_binary::kMatrixWeightPrecisionBits * 3);
E copy{E::Work, 4000};
copy.workClass = E::LocalCopy;
abstract[0] = {copy, {E::Load, 1}};
auto charged = analyzePimExecution(abstract, target);
PimPeriodicOptions sensitivity;
sensitivity.ignoreLocalCopies = true;
auto omitted = analyzePimExecution(abstract, target, sensitivity);
TEST_CHECK(*charged.getObject("memory_controller")->getNumber("interval_ns")
- *omitted.getObject("memory_controller")->getNumber("interval_ns") == 4000);
}
static void testAccountingAndOverflow(PimInstructionTiming target) {
auto record = [](O op, int64_t bytes) {
pim_binary::InstructionRecord r;
r.opcode = op; r.generic3 = bytes;
return r;
};
PimPeriodicOptions options;
options.scalarIssue = false;
options.workBreakdown = true;
auto setbw = record(O::setbw, 0);
setbw.generic1 = setbw.generic2 = 8;
auto matrixInstruction = record(O::mvmul, 0);
matrixInstruction.r2OrImm = pim_binary::kMatrixWeightPrecisionBits;
PimInstructionProgram reference;
reference.instructions = {binary({setbw, matrixInstruction, record(O::vvadd, 65), record(O::vmv, 3), record(O::lmv, 65), record(O::lldi, 65)})};
reference.arrayGroups = {{3}};
E matrix{E::Work, getPimMatrixWorkNs(target, 8,
getPimMatrixOutputBytes(target, 2, pim_binary::kMatrixWeightPrecisionBits, 3, 1))};
matrix.workClass = E::Matrix;
E vector{E::Work, getPimVectorWorkNs(target, 65, 1, 1, true)};
vector.workClass = E::Vector;
E vectorCopy{E::Work, getPimVectorWorkNs(target, 3, 1, 1, false)};
vectorCopy.workClass = E::VectorCopy;
E copy{E::Work, target.getLocalMemoryServiceNs(65, false)
+ target.getLocalMemoryServiceNs(65, true) + target.processorPeriodNs};
copy.workClass = E::LocalCopy;
E fill{E::Work, target.getLocalMemoryServiceNs(65, true) + target.processorPeriodNs};
fill.workClass = E::Fill;
auto exact = analyzePimInstructions(reference, target, 2, options);
auto abstract = analyzePimExecution({{matrix, vector, vectorCopy, copy, fill}}, target, options);
TEST_CHECK(*abstract.getArray("resource_service_ns") == *exact.getArray("resource_service_ns"));
TEST_CHECK(abstract.getObject("memory_controller")->getNumber("interval_ns")
== exact.getObject("memory_controller")->getNumber("interval_ns"));
auto repeated = analyzePimExecution({{{E::Repeat, 7, 0, 0, 0, {matrix, vector, vectorCopy, copy, fill}}}}, target, options);
for (size_t i = 0; i < 7; ++i)
TEST_CHECK(*(*repeated.getArray("resource_service_ns"))[i].getAsInteger()
== 7 * *(*abstract.getArray("resource_service_ns"))[i].getAsInteger());
auto nested = analyzePimExecution({{{E::Repeat, 1000000, 0, 0, 0,
{{E::Repeat, 1000000, 0, 0, 0, {fill}}}}}}, target, options);
TEST_CHECK(*(*nested.getArray("resource_service_ns"))[2].getAsInteger()
== static_cast<int64_t>(1000000000000ull * fill.amount));
TEST_CHECK(nested.getInteger("expanded_graph_events").value() < 10);
PimExecutionProgram pattern(2), expanded(2);
pattern[0] = {{E::Repeat, 3, 0, 0, 0, {{E::Load, 65},
{E::Repeat, 5, 0, 0, 0, {matrix, vector, copy, fill}}}}};
for (size_t repeat = 0; repeat < 3; ++repeat) {
expanded[0].push_back({E::Load, 65});
for (size_t work = 0; work < 5; ++work)
expanded[0].insert(expanded[0].end(), {matrix, vector, copy, fill});
}
pattern[1] = expanded[1] = {{E::Work, 311}, {E::Load, 257}};
auto compactCost = analyzePimExecution(pattern, target, options);
auto expandedCost = analyzePimExecution(expanded, target, options);
for (const char *policy : {"interval_ns", "alternative_tie_interval_ns"})
TEST_CHECK(compactCost.getObject("memory_controller")->getNumber(policy)
== expandedCost.getObject("memory_controller")->getNumber(policy));
TEST_CHECK(*compactCost.getArray("resource_service_ns") == *expandedCost.getArray("resource_service_ns"));
std::vector<E> body{{E::Load, 1}};
for (size_t i = 0; i < 10000; ++i) {
E event{E::Work, 1}; event.workClass = i % 2 ? E::Vector : E::Fill;
body.push_back(event);
}
auto cachedWork = analyzePimExecution({{{E::Repeat, 10000, 0, 0, 0, body}}}, target, options);
TEST_CHECK(cachedWork.getInteger("global_read_requests") == 10000);
TEST_CHECK(cachedWork.getInteger("expanded_graph_events").value() < 100000);
TEST_CHECK((*cachedWork.getArray("resource_service_ns"))[1].getAsInteger() == 50000000);
TEST_CHECK((*cachedWork.getArray("resource_service_ns"))[2].getAsInteger() == 50000000);
uint64_t maximum = std::numeric_limits<uint64_t>::max();
for (auto program : std::vector<PimExecutionProgram>{
{{{E::Work, maximum}}}, {{{E::Load, maximum}}},
{{{E::Repeat, maximum, 0, 0, 0, {fill}}}},
{{{E::Work, uint64_t(1) << 62}, {E::Work, uint64_t(1) << 62}}}})
TEST_CHECK(analyzePimExecution(program, target).getString("status") == "unavailable");
target.sharedMemoryReadLatencyNs = maximum;
TEST_CHECK(analyzePimExecution({{{E::Load, 65}}}, target).getString("status") == "unavailable");
}
static PimPeriodicOptions spatialDetails() {
PimPeriodicOptions options;
options.workBreakdown = true;
return options;
}
static void testCoreAssignments() {
MLIRContext context;
context.loadDialect<func::FuncDialect, spatial::SpatialDialect>();
OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
OpBuilder builder(module->getBodyRegion());
auto loc = module->getLoc();
auto function = func::FuncOp::create(builder, loc, "test", builder.getFunctionType({}, {}));
builder.setInsertionPointToStart(function.addEntryBlock());
auto empty = [&](std::optional<int32_t> id) {
auto compute = spatial::SpatScheduledCompute::create(builder, loc, TypeRange{}, ValueRange{}, ValueRange{});
if (id) compute->setAttr("coreId", builder.getI32IntegerAttr(*id));
compute.getBody().push_back(new Block);
builder.setInsertionPointToStart(&compute.getBody().front());
spatial::SpatYieldOp::create(builder, loc, ValueRange{});
builder.setInsertionPointAfter(compute);
return compute;
};
auto first = empty(0);
auto second = empty(1);
func::ReturnOp::create(builder, loc);
TEST_CHECK(succeeded(verify(*module)));
auto target = hardware(1);
auto reason = [&]() {
auto report = analyzePimSpatialSchedule(*module, target);
return report.getString("reason").value_or("").str();
};
TEST_CHECK(reason().find("physical core") == std::string::npos);
second->setAttr("coreId", builder.getI32IntegerAttr(0));
TEST_CHECK(reason() == "duplicate physical core assignment: 0");
TEST_CHECK(evaluatePimSpatialSchedule(*module, target).cost.loweringSupport
== spatial::StaticLoweringSupport::Unsupported);
second->removeAttr("coreId");
TEST_CHECK(reason() == "missing physical core");
TEST_CHECK(evaluatePimSpatialSchedule(*module, target).cost.loweringSupport
== spatial::StaticLoweringSupport::Unknown);
for (int32_t id : {-1, 2}) {
second->setAttr("coreId", builder.getI32IntegerAttr(id));
TEST_CHECK(reason() == "physical core ID is outside the target: " + std::to_string(id));
}
second.erase();
builder.setInsertionPoint(function.getBody().front().getTerminator());
auto batch = spatial::SpatScheduledComputeBatch::create(builder, loc, TypeRange{},
builder.getI32IntegerAttr(2), ValueRange{}, ValueRange{});
batch.getBody().push_back(new Block);
batch.getBody().front().addArgument(builder.getIndexType(), loc);
builder.setInsertionPointToStart(&batch.getBody().front());
spatial::SpatYieldOp::create(builder, loc, ValueRange{});
TEST_CHECK(reason() == "missing physical batch cores");
batch->setAttr("coreIds", builder.getDenseI32ArrayAttr({1, 1}));
TEST_CHECK(reason() == "duplicate physical core assignment: 1");
batch->setAttr("coreIds", builder.getDenseI32ArrayAttr({1, 0}));
TEST_CHECK(reason() == "duplicate physical core assignment: 0");
TEST_CHECK(evaluatePimSpatialSchedule(*module, target).cost.loweringSupport
== spatial::StaticLoweringSupport::Unsupported);
first.erase();
batch->setAttr("coreIds", builder.getDenseI32ArrayAttr({1, 2}));
TEST_CHECK(reason() == "physical core ID is outside the target: 2");
}
static void testDependencyTokens() {
PimDependencyTokens fresh, drained;
TEST_CHECK(fresh.empty() && fresh.remaining().empty());
drained.initialize(1);
TEST_CHECK(drained.consume() == 0);
TEST_CHECK(drained.empty() && drained.remaining().empty());
// Empty queues have identical future behavior, regardless of their capacity.
for (uint64_t time : {0, 7, 7, 123}) {
fresh.push(time); drained.push(time);
TEST_CHECK(fresh.remaining() == drained.remaining());
TEST_CHECK(fresh.consume() == drained.consume());
TEST_CHECK(fresh.empty() && drained.empty());
}
// Consumed prefixes and compaction cannot change FIFO contents or readiness.
drained.initialize(3);
TEST_CHECK(drained.consume() == 0);
drained.push(11);
TEST_CHECK(drained.consume() == 0);
TEST_CHECK(drained.consume() == 0);
drained.push(11);
drained.push(12);
for (uint64_t time : {11, 11, 12}) TEST_CHECK(drained.consume() == time);
TEST_CHECK(drained.empty());
drained.initialize(0);
TEST_CHECK(drained.empty() && drained.remaining().empty());
}
static void testReadOnlySpatial(PimInstructionTiming target, bool batched) {
MLIRContext context;
context.loadDialect<arith::ArithDialect, func::FuncDialect, scf::SCFDialect,
tensor::TensorDialect, linalg::LinalgDialect, spatial::SpatialDialect>();
auto module = ModuleOp::create(UnknownLoc::get(&context));
OpBuilder builder(module.getBodyRegion());
auto loc = module.getLoc();
auto type = RankedTensorType::get({1, 64}, builder.getF32Type());
auto function = func::FuncOp::create(builder, loc, "test", builder.getFunctionType({}, {}));
builder.setInsertionPointToStart(function.addEntryBlock());
Value constant = arith::ConstantOp::create(builder, loc,
DenseElementsAttr::get(type, builder.getF32FloatAttr(2)));
Operation *compute;
if (batched) {
compute = spatial::SpatScheduledComputeBatch::create(builder, loc, TypeRange{},
builder.getI32IntegerAttr(2), ValueRange{}, ValueRange{constant});
compute->setAttr("coreIds", builder.getDenseI32ArrayAttr({0, 1}));
} else {
compute = spatial::SpatScheduledCompute::create(builder, loc, TypeRange{}, ValueRange{}, ValueRange{constant});
compute->setAttr("coreId", builder.getI32IntegerAttr(0));
}
Block *body = new Block;
compute->getRegion(0).push_back(body);
if (batched) body->addArgument(builder.getIndexType(), loc);
Value input = body->addArgument(type, loc);
builder.setInsertionPointToStart(body);
Value zero = arith::ConstantIndexOp::create(builder, loc, 0);
Value upper = arith::ConstantIndexOp::create(builder, loc, 28);
Value one = arith::ConstantIndexOp::create(builder, loc, 1);
Value local = tensor::EmptyOp::create(builder, loc, type.getShape(), type.getElementType());
auto loop = scf::ForOp::create(builder, loc, zero, upper, one);
builder.setInsertionPointToStart(loop.getBody());
auto add = spatial::SpatVAddOp::create(builder, loc, type, local, input);
builder.setInsertionPointAfter(loop);
spatial::SpatWaitOp::create(builder, loc, zero, zero);
spatial::SpatYieldOp::create(builder, loc, ValueRange{});
builder.setInsertionPointAfter(compute);
func::ReturnOp::create(builder, loc);
TEST_CHECK(succeeded(verify(module)));
std::string before, after;
llvm::raw_string_ostream os(before); module.print(os);
auto report = analyzePimSpatialSchedule(module, target, spatialDetails());
auto fast = analyzePimSpatialSchedule(module, target);
PimPeriodicOptions proofOptions;
proofOptions.verifyPeriodicState = true;
auto typed = evaluatePimSpatialSchedule(module, target, proofOptions);
TEST_CHECK(typed.cost.available() && typed.cost.convergenceProved() == true);
TEST_CHECK(typed.cost.loweringSupport == spatial::StaticLoweringSupport::Supported);
TEST_CHECK(typed.report.getBoolean("lowering_supported") == typed.cost.loweringSupported());
TEST_CHECK(typed.report.getBoolean("convergence_proved") == typed.cost.convergenceProved());
TEST_CHECK(!fast.getArray("request_order") && !fast.getArray("resource_service_ns"));
TEST_CHECK(!fast.getObject("without_local_copy_work") && !fast.getArray("work_class_ns"));
for (const char *field : {"interval_ns", "alternative_tie_interval_ns"})
TEST_CHECK(fast.getObject("memory_controller")->getNumber(field)
== report.getObject("memory_controller")->getNumber(field));
TEST_CHECK(report.getBoolean("safe_for_selection") == false);
TEST_CHECK(report.getInteger("global_read_requests") == (batched ? 2 : 28));
TEST_CHECK(report.getInteger("global_read_bytes") == (batched ? 2 : 28) * 64);
TEST_CHECK(report.getInteger("constant_read_requests") == (batched ? 2 : 28));
llvm::raw_string_ostream out(after); module.print(out);
TEST_CHECK(before == after);
add->setAttr("test.annotation", builder.getUnitAttr());
report = analyzePimSpatialSchedule(module, target, spatialDetails());
TEST_CHECK(report.getInteger("global_read_requests") == (batched ? 2 : 28));
builder.setInsertionPoint(body->getTerminator());
Value source = tensor::EmptyOp::create(builder, loc, ArrayRef<int64_t>{1, 3, 1, 224}, builder.getF32Type());
Value destination = tensor::EmptyOp::create(builder, loc, ArrayRef<int64_t>{1, 1, 224, 3}, builder.getF32Type());
linalg::TransposeOp::create(builder, loc, source, destination, ArrayRef<int64_t>{0, 2, 3, 1});
TEST_CHECK(succeeded(verify(module)));
report = analyzePimSpatialSchedule(module, target, spatialDetails());
auto expectedCopyWork = 224 * getPimVectorWorkNs(target, 3, 1, 1, false) * (batched ? 2 : 1);
auto *classes = report.getArray("work_class_ns");
auto *columns = report.getArray("work_class_columns");
TEST_CHECK(classes && columns && classes->size() == 6 && columns->size() == 6);
const char *names[] = {"matrix", "vector", "local_copy", "fill", "vector_copy", "scalar_issue"};
uint64_t totalWork = 0;
for (size_t i = 0; i < classes->size(); ++i) {
TEST_CHECK((*columns)[i].getAsString() == names[i]);
totalWork += *(*classes)[i].getAsInteger();
}
TEST_CHECK((*classes)[E::VectorCopy].getAsInteger() == expectedCopyWork);
TEST_CHECK((*classes)[E::LocalCopy].getAsInteger() == 0);
TEST_CHECK((*classes)[E::ScalarIssue].getAsInteger() == 0);
TEST_CHECK(totalWork == report.getInteger("compute_and_local_work_ns"));
auto softmax = spatial::SpatSoftmaxOp::create(builder, loc, type, input);
TEST_CHECK(succeeded(verify(module)));
TEST_CHECK(analyzePimSpatialSchedule(module, target, spatialDetails()).getString("status") == "unavailable");
softmax.erase();
source = tensor::EmptyOp::create(builder, loc, ArrayRef<int64_t>{1, 0, 1, 224}, builder.getF32Type());
destination = tensor::EmptyOp::create(builder, loc, ArrayRef<int64_t>{1, 1, 224, 0}, builder.getF32Type());
linalg::TransposeOp::create(builder, loc, source, destination, ArrayRef<int64_t>{0, 2, 3, 1});
TEST_CHECK(succeeded(verify(module)));
report = analyzePimSpatialSchedule(module, target, spatialDetails());
TEST_CHECK(report.getString("reason") == "zero-element transpose has no abstract timing contract");
module.erase();
}
static void testHostSlice(PimInstructionTiming target) {
MLIRContext context;
context.loadDialect<arith::ArithDialect, func::FuncDialect, scf::SCFDialect,
tensor::TensorDialect, spatial::SpatialDialect>();
auto module = ModuleOp::create(UnknownLoc::get(&context));
OpBuilder builder(module.getBodyRegion());
auto loc = module.getLoc();
auto type = RankedTensorType::get({1, 3, 4, 4}, builder.getF32Type());
auto function = func::FuncOp::create(builder, loc, "test", builder.getFunctionType({type}, {}));
builder.setInsertionPointToStart(function.addEntryBlock());
auto batch = spatial::SpatScheduledComputeBatch::create(builder, loc, TypeRange{},
builder.getI32IntegerAttr(2), ValueRange{}, ValueRange{function.getArgument(0)});
batch->setAttr("coreIds", builder.getDenseI32ArrayAttr({0, 1}));
Block *body = new Block;
batch.getBody().push_back(body);
body->addArgument(builder.getIndexType(), loc);
Value input = body->addArgument(type, loc);
builder.setInsertionPointToStart(body);
auto offsets = getAsOpFoldResult(builder.getI64ArrayAttr({0, 0, 1, 0}));
auto sizes = getAsOpFoldResult(builder.getI64ArrayAttr({1, 3, 1, 4}));
auto strides = getAsOpFoldResult(builder.getI64ArrayAttr({1, 1, 1, 1}));
auto slice = tensor::ExtractSliceOp::create(builder, loc, input, offsets, sizes, strides);
spatial::SpatReluOp::create(builder, loc, slice.getType(), slice.getResult());
spatial::SpatYieldOp::create(builder, loc, ValueRange{});
builder.setInsertionPointAfter(batch);
func::ReturnOp::create(builder, loc);
TEST_CHECK(succeeded(verify(module)));
auto report = analyzePimSpatialSchedule(module, target, spatialDetails());
TEST_CHECK(report.getInteger("global_read_requests") == 6);
TEST_CHECK(report.getInteger("global_read_bytes") == 24);
module.erase();
}
static void testCapturedVectorWeights(PimInstructionTiming target) {
MLIRContext context;
context.loadDialect<arith::ArithDialect, func::FuncDialect, tensor::TensorDialect,
spatial::SpatialDialect>();
auto module = ModuleOp::create(UnknownLoc::get(&context));
OpBuilder builder(module.getBodyRegion());
auto loc = module.getLoc();
auto type = RankedTensorType::get({1, 64}, builder.getF32Type());
auto function = func::FuncOp::create(builder, loc, "test", builder.getFunctionType({}, {}));
builder.setInsertionPointToStart(function.addEntryBlock());
Value zero = arith::ConstantOp::create(builder, loc,
DenseElementsAttr::get(type, builder.getF32FloatAttr(0)));
auto compute = spatial::SpatScheduledCompute::create(builder, loc, TypeRange{},
ValueRange{zero}, ValueRange{});
compute->setAttr("coreId", builder.getI32IntegerAttr(0));
Block *body = new Block;
compute.getBody().push_back(body);
Value weight = body->addArgument(type, loc);
builder.setInsertionPointToStart(body);
Value local = tensor::EmptyOp::create(builder, loc, type.getShape(), type.getElementType());
spatial::SpatVAddOp::create(builder, loc, type, local, weight);
spatial::SpatYieldOp::create(builder, loc, ValueRange{});
builder.setInsertionPointAfter(compute);
func::ReturnOp::create(builder, loc);
TEST_CHECK(succeeded(verify(module)));
auto report = analyzePimSpatialSchedule(module, target, spatialDetails());
// Captured weight operands stay block arguments; codegen cannot fold this load to LLDI.
TEST_CHECK(report.getInteger("global_read_requests") == 1);
TEST_CHECK(report.getInteger("constant_read_bytes") == 64);
module.erase();
}
static void testPublicationAndOutput(PimInstructionTiming target) {
MLIRContext context;
context.loadDialect<arith::ArithDialect, func::FuncDialect, tensor::TensorDialect,
spatial::SpatialDialect>();
auto module = ModuleOp::create(UnknownLoc::get(&context));
OpBuilder builder(module.getBodyRegion());
auto loc = module.getLoc();
auto inputType = RankedTensorType::get({1, 2, 3, 5}, builder.getF32Type());
auto outputType = RankedTensorType::get({1, 4, 3, 5}, builder.getF32Type());
auto function = func::FuncOp::create(builder, loc, "test", builder.getFunctionType({}, {outputType}));
builder.setInsertionPointToStart(function.addEntryBlock());
auto compute = spatial::SpatScheduledCompute::create(builder, loc, TypeRange{outputType},
ValueRange{}, ValueRange{});
compute->setAttr("coreId", builder.getI32IntegerAttr(0));
compute.getBody().push_back(new Block);
builder.setInsertionPointToStart(&compute.getBody().front());
Value zero = arith::ConstantIndexOp::create(builder, loc, 0);
Value one = arith::ConstantIndexOp::create(builder, loc, 1);
Value input = tensor::EmptyOp::create(builder, loc, inputType.getShape(), inputType.getElementType());
auto concat = spatial::SpatConcatOp::create(builder, loc, outputType,
builder.getI64IntegerAttr(1), ValueRange{input, input});
auto slice = tensor::ExtractSliceOp::create(builder, loc, concat.getOutput(),
getAsOpFoldResult(builder.getI64ArrayAttr({0, 0, 1, 0})),
getAsOpFoldResult(builder.getI64ArrayAttr({1, 4, 1, 5})),
getAsOpFoldResult(builder.getI64ArrayAttr({1, 1, 1, 1})));
spatial::SpatHostStoreSyncOp::create(builder, loc, zero, one, zero, zero, slice.getResult());
spatial::SpatWaitOp::create(builder, loc, one, one);
spatial::SpatYieldOp::create(builder, loc, ValueRange{concat.getOutput()});
builder.setInsertionPointAfter(compute);
auto reader = spatial::SpatScheduledCompute::create(builder, loc, TypeRange{}, ValueRange{}, ValueRange{});
reader->setAttr("coreId", builder.getI32IntegerAttr(1));
reader.getBody().push_back(new Block);
builder.setInsertionPointToStart(&reader.getBody().front());
zero = arith::ConstantIndexOp::create(builder, loc, 0);
one = arith::ConstantIndexOp::create(builder, loc, 1);
spatial::SpatWaitOp::create(builder, loc, zero, one);
spatial::SpatSyncOp::create(builder, loc, zero, one);
spatial::SpatYieldOp::create(builder, loc, ValueRange{});
builder.setInsertionPointAfter(reader);
func::ReturnOp::create(builder, loc, compute.getResults());
TEST_CHECK(succeeded(verify(module)));
auto report = analyzePimSpatialSchedule(module, target, spatialDetails());
TEST_CHECK(report.getInteger("global_write_requests") == 5);
TEST_CHECK(report.getInteger("global_write_bytes") == 80);
auto *runs = (*report.getArray("request_order"))[0].getAsObject()->getArray("runs");
TEST_CHECK(runs->size() == 2);
TEST_CHECK((*(*runs)[0].getAsArray())[1].getAsInteger() == 60);
TEST_CHECK((*(*runs)[1].getAsArray())[1].getAsInteger() == 5);
TEST_CHECK((*(*runs)[1].getAsArray())[2].getAsInteger() == 4);
module.erase();
}
static void testUnrepresentableSpatial(PimInstructionTiming target) {
MLIRContext context;
context.loadDialect<arith::ArithDialect, func::FuncDialect, tensor::TensorDialect,
scf::SCFDialect, spatial::SpatialDialect>();
for (bool loopOverflow : {false, true}) {
auto module = ModuleOp::create(UnknownLoc::get(&context));
OpBuilder builder(module.getBodyRegion());
auto loc = module.getLoc();
auto function = func::FuncOp::create(builder, loc, "test", builder.getFunctionType({}, {}));
builder.setInsertionPointToStart(function.addEntryBlock());
auto compute = spatial::SpatScheduledCompute::create(builder, loc, TypeRange{}, ValueRange{}, ValueRange{});
compute->setAttr("coreId", builder.getI32IntegerAttr(0));
compute.getBody().push_back(new Block);
builder.setInsertionPointToStart(&compute.getBody().front());
if (loopOverflow) {
Value lower = arith::ConstantIndexOp::create(builder, loc, std::numeric_limits<int64_t>::min());
Value upper = arith::ConstantIndexOp::create(builder, loc, std::numeric_limits<int64_t>::max());
Value step = arith::ConstantIndexOp::create(builder, loc, 1);
scf::ForOp::create(builder, loc, lower, upper, step);
} else {
auto type = RankedTensorType::get({std::numeric_limits<int64_t>::max(), 2}, builder.getF32Type());
Value input = tensor::EmptyOp::create(builder, loc, type.getShape(), type.getElementType());
spatial::SpatReluOp::create(builder, loc, type, input);
}
spatial::SpatYieldOp::create(builder, loc, ValueRange{});
builder.setInsertionPointAfter(compute);
func::ReturnOp::create(builder, loc);
TEST_CHECK(succeeded(verify(module)));
auto report = analyzePimSpatialSchedule(module, target, spatialDetails());
TEST_CHECK(report.getString("status") == "unavailable");
TEST_CHECK(!report.getBoolean("safe_for_selection").value());
module.erase();
}
}
static void testMatrixGeometry(PimInstructionTiming target) {
MLIRContext context;
context.loadDialect<arith::ArithDialect, func::FuncDialect, tensor::TensorDialect,
spatial::SpatialDialect>();
target.matrixCellBits = 2;
for (bool encoded : {false, true})
for (int64_t rows : {64, 128, 192})
for (int64_t columns : {64, 128, 192, 256}) {
auto module = ModuleOp::create(UnknownLoc::get(&context));
OpBuilder builder(module.getBodyRegion());
auto loc = module.getLoc();
auto inputType = RankedTensorType::get({1, rows}, builder.getF32Type());
auto weightType = RankedTensorType::get({rows, columns}, builder.getF32Type(),
encoded ? builder.getStringAttr("unsupported-layout") : Attribute{});
auto resultType = RankedTensorType::get({1, columns}, builder.getF32Type());
auto function = func::FuncOp::create(builder, loc, "test", builder.getFunctionType({}, {}));
builder.setInsertionPointToStart(function.addEntryBlock());
Value weight = arith::ConstantOp::create(builder, loc,
DenseElementsAttr::get(weightType, builder.getF32FloatAttr(1)));
auto compute = spatial::SpatScheduledCompute::create(builder, loc, TypeRange{},
ValueRange{weight}, ValueRange{});
compute->setAttr("coreId", builder.getI32IntegerAttr(0));
compute.getBody().push_back(new Block);
Value captured = compute.getBody().front().addArgument(weightType, loc);
builder.setInsertionPointToStart(&compute.getBody().front());
Value input = tensor::EmptyOp::create(builder, loc, inputType.getShape(), inputType.getElementType());
spatial::SpatVMMOp::create(builder, loc, resultType, captured, input);
spatial::SpatYieldOp::create(builder, loc, ValueRange{});
builder.setInsertionPointAfter(compute);
func::ReturnOp::create(builder, loc);
TEST_CHECK(succeeded(verify(module)));
auto report = analyzePimSpatialSchedule(module, target);
if (encoded || columns % target.matrixColumns || static_cast<uint64_t>(rows) > target.matrixRows)
TEST_CHECK(report.getString("reason") == "unavailable physical matrix weight geometry");
else TEST_CHECK(report.getString("status") != "unavailable");
module.erase();
}
}
static void testOutputMultiplicity(PimInstructionTiming target) {
MLIRContext context;
context.loadDialect<arith::ArithDialect, func::FuncDialect, tensor::TensorDialect,
spatial::SpatialDialect>();
for (bool repeatedResult : {false, true}) {
auto module = ModuleOp::create(UnknownLoc::get(&context));
OpBuilder builder(module.getBodyRegion());
auto loc = module.getLoc();
auto type = RankedTensorType::get({1, 64}, builder.getF32Type());
auto function = func::FuncOp::create(builder, loc, "test", builder.getFunctionType({}, {type, type}));
builder.setInsertionPointToStart(function.addEntryBlock());
SmallVector<Type> results(repeatedResult ? 1 : 2, type);
auto compute = spatial::SpatScheduledCompute::create(builder, loc, results, ValueRange{}, ValueRange{});
compute->setAttr("coreId", builder.getI32IntegerAttr(0));
compute.getBody().push_back(new Block);
builder.setInsertionPointToStart(&compute.getBody().front());
Value input = tensor::EmptyOp::create(builder, loc, type.getShape(), type.getElementType());
Value output = spatial::SpatReluOp::create(builder, loc, type, input);
SmallVector<Value> yielded(results.size(), output);
spatial::SpatYieldOp::create(builder, loc, yielded);
builder.setInsertionPointAfter(compute);
func::ReturnOp::create(builder, loc, ValueRange{compute.getResult(0), compute.getResult(repeatedResult ? 0 : 1)});
TEST_CHECK(succeeded(verify(module)));
auto report = analyzePimSpatialSchedule(module, target, spatialDetails());
if (repeatedResult)
TEST_CHECK(report.getString("reason") == "unavailable indirect compute output collection");
else {
TEST_CHECK(report.getInteger("global_write_requests") == 2);
TEST_CHECK(report.getInteger("global_write_bytes") == 128);
auto *runs = (*report.getArray("request_order"))[0].getAsObject()->getArray("runs");
TEST_CHECK((*(*runs)[0].getAsArray())[2].getAsInteger() == 2);
}
module.erase();
}
}
static void testPeriodicStateProof() {
auto target = hardware(1);
target.processorCount = 3;
target.sharedMemoryReadReturnsPayload = false;
target.localMemoryPeriodNs = 1;
target.localMemoryReadLatencyCycles = 20;
target.localMemoryWriteLatencyCycles = 30;
target.transferWidthBytes = target.sharedMemoryWidthBytes = 32;
target.sharedMemoryReadLatencyNs = 40;
target.sharedMemoryWriteLatencyNs = 50;
target.sharedMemoryRequestLatencyNs = target.sharedMemoryResponseLatencyNs = {6, 6, 6};
target.interProcessorLatencyNs = {0, 3, 4, 3, 0, 3, 4, 3, 0};
PimExecutionProgram program{
{{E::Work, 377}, {E::Store, 875}, {E::Work, 545}, {E::Load, 367},
{E::Work, 333}, {E::Store, 80}, {E::Signal, 0, 1}, {E::Wait, 0, 0, 0, 1}},
{{E::Work, 924}, {E::Load, 312}, {E::Work, 70}, {E::Store, 434},
{E::Work, 337}, {E::Load, 275}, {E::Signal, 0, 2}, {E::Wait, 0, 0, 0, 1}},
{{E::Work, 158}, {E::Store, 565}, {E::Work, 126}, {E::Load, 245},
{E::Work, 1124}, {E::Store, 367}, {E::Signal, 0, 0}, {E::Wait, 0, 0, 0, 1}}};
auto baseline = analyzePimExecution(program, target);
PimPeriodicOptions options;
options.verifyPeriodicState = true;
options.workBreakdown = true;
auto verified = analyzePimExecution(program, target, options);
auto *queue = verified.getObject("memory_controller");
auto *legacy = queue->getObject("legacy_interval_window");
TEST_CHECK(legacy);
for (const char *field : {"interval_ns", "alternative_tie_interval_ns"})
TEST_CHECK(legacy->getNumber(field) == baseline.getObject("memory_controller")->getNumber(field));
auto *state = queue->getObject("periodic_state");
TEST_CHECK(state && state->getString("status") == "proved");
TEST_CHECK(state->getInteger("period_generations") == 3);
TEST_CHECK(state->getNumber("interval_ns") == 17485.0 / 3);
TEST_CHECK(queue->getNumber("interval_ns") == state->getNumber("interval_ns"));
TEST_CHECK(legacy->getNumber("interval_ns") != state->getNumber("interval_ns"));
// A fast writer builds signal tokens. Repeated local intervals are insufficient.
program = {{{E::Signal, 0, 1}, {E::Work, 10}},
{{E::Wait, 0, 0, 0, 1}, {E::Work, 100}}};
verified = analyzePimExecution(program, target, options);
state = verified.getObject("memory_controller")->getObject("periodic_state");
TEST_CHECK(state && state->getString("status") == "unproved");
}
static void testCompleteSchedulerState() {
auto target = hardware(1);
PimPeriodicOptions options;
options.verifyPeriodicState = true;
auto verify = [&](const PimExecutionProgram &program, bool expected) {
auto baseline = analyzePimExecution(program, target);
auto typed = evaluatePimExecution(program, target, options);
TEST_CHECK(typed.cost.convergenceProved() == expected);
auto detailed = std::move(typed.report);
auto *queue = detailed.getObject("memory_controller");
TEST_CHECK(queue);
TEST_CHECK(detailed.getString("status") == (typed.cost.available() ? "estimated" : "unavailable"));
if (!typed.cost.available()) TEST_CHECK(detailed.getString("reason") == typed.cost.reason);
if (expected) {
TEST_CHECK(queue->getNumber("interval_ns") == typed.cost.primary);
TEST_CHECK(queue->getNumber("alternative_tie_interval_ns") == typed.cost.alternative);
} else {
for (const char *field : {"interval_ns", "alternative_tie_interval_ns"})
TEST_CHECK(queue->getNumber(field) == baseline.getObject("memory_controller")->getNumber(field));
}
size_t policy = 0;
for (const char *field : {"periodic_state", "alternative_periodic_state"}) {
auto *state = queue->getObject(field);
TEST_CHECK(state && (state->getString("status") == "proved") == expected);
const auto &evidence = typed.cost.recurrence[policy++];
TEST_CHECK(evidence.proved() == expected);
TEST_CHECK(evidence.intervalNs == state->getNumber("interval_ns"));
if (expected) TEST_CHECK(state->getBoolean("complete_scheduler_state") == true);
}
return detailed;
};
verify({{{E::Work, 0}}}, false);
auto steady = verify({{{E::Work, 100}}}, true);
auto *state = steady.getObject("memory_controller")->getObject("periodic_state");
TEST_CHECK(state->getInteger("period_generations") == 1);
TEST_CHECK(state->getNumber("interval_ns") == 100);
verify({{{E::Work, uint64_t{1} << 60}}}, true);
// Simultaneous signals, drained token queues, and controller arrivals must
// recur under each deterministic arbitration policy.
verify({{{E::Work, 10}, {E::Signal, 0, 1}, {E::Wait, 0, 0, 0, 1}},
{{E::Work, 10}, {E::Signal, 0, 0}, {E::Wait, 0, 0, 0, 1}}}, true);
verify({{{E::Load, 65}, {E::Signal, 0, 1}, {E::Wait, 0, 0, 0, 1}},
{{E::Load, 65}, {E::Signal, 0, 0}, {E::Wait, 0, 0, 0, 1}}}, true);
// Repeated output intervals do not prove uniform progress on every core.
auto unequal = verify({{{E::Work, 7}}, {{E::Work, 11}}}, false);
TEST_CHECK(unequal.getObject("memory_controller")->getString("status") == "estimated");
// The complete state retains growing dependency-token queues.
auto backlog = verify({{{E::Signal, 0, 1}, {E::Work, 10}},
{{E::Wait, 0, 0, 0, 1}, {E::Work, 100}}}, false);
TEST_CHECK(backlog.getObject("memory_controller")->getString("status") == "estimated");
// Both endpoints exist, but neither WAIT can receive its signal.
auto deadlock = verify({{{E::Wait, 0, 0, 0, 1}, {E::Signal, 0, 1}},
{{E::Wait, 0, 0, 0, 1}, {E::Signal, 0, 0}}}, false);
TEST_CHECK(deadlock.getObject("memory_controller")->getString("status") == "unavailable");
verify({{{E::Send, 65, 1}, {E::Receive, 65, 1}},
{{E::Send, 65, 0}, {E::Receive, 65, 0}}}, false);
// An independent active core must not hide an inactive dependency cycle.
verify({{{E::Work, 10}}, {{E::Wait, 0, 0, 0, 1}, {E::Signal, 0, 1}}}, false);
auto overflow = analyzePimExecution({{{E::Work, uint64_t{1} << 62}}}, target, options);
TEST_CHECK(overflow.getString("status") == "unavailable");
}
int main() {
testDependencyTokens();
testCoreAssignments();
testPeriodicStateProof();
testCompleteSchedulerState();
for (int64_t period : {1, 5}) {
testLateRecurrenceAuthority(hardware(period));
testTypedResult(hardware(period));
testSharedEngine(hardware(period));
testAccountingAndOverflow(hardware(period));
testMatrixGeometry(hardware(period));
testOutputMultiplicity(hardware(period));
testHostSlice(hardware(period));
testCapturedVectorWeights(hardware(period));
testPublicationAndOutput(hardware(period));
testUnrepresentableSpatial(hardware(period));
testReadOnlySpatial(hardware(period), false);
testReadOnlySpatial(hardware(period), true);
}
}
test/PIM/PimPipelineSelectionTest.cpp
+284 -40
View File
@@ -1,11 +1,259 @@
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ScheduleEvaluation.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.hpp"
#include "llvm/Support/JSON.h"
#include <cassert>
#include "TestCheck.hpp"
#include <limits>
#include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
#include "src/Accelerators/PIM/Passes/PIMPasses.h"
#include "src/Accelerators/PIM/Compiler/PimPipelineSelection.hpp"
#include "src/Accelerators/PIM/Compiler/PimCodeGen.hpp"
#include "src/Accelerators/PIM/Compiler/PimCompilerUtils.hpp"
#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
#include "src/Accelerators/PIM/Common/Support/FileSystemUtils.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTargetResources.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/Passes/Transforms/Bufferization/OpBufferizationInterfaces.hpp"
#include "mlir/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/IR/Verifier.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/MemoryBuffer.h"
using namespace onnx_mlir;
using namespace onnx_mlir::spatial;
static WholeScheduleCost scheduleCost(double interval) {
WholeScheduleCost cost;
cost.primary = cost.alternative = interval;
cost.selectionMargin = 0.021;
cost.loweringSupport = StaticLoweringSupport::Supported;
for (auto &recurrence : cost.recurrence)
recurrence = {ScheduleRecurrenceEvidence::Proved, 1, interval, {}};
return cost;
}
static void testEmptyGraphPreparation() {
using namespace mlir;
MLIRContext context;
context.loadDialect<func::FuncDialect, SpatialDialect>();
OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
OpBuilder builder(module->getBodyRegion());
auto function = func::FuncOp::create(builder, module->getLoc(), "main_graph",
builder.getFunctionType({}, {}));
builder.setInsertionPointToStart(function.addEntryBlock());
func::ReturnOp::create(builder, module->getLoc());
SchedulingTarget target;
target.processorCount = target.residentWeightCapacity = 1;
target.interProcessorLatencyNs = {0};
for (bool collectCandidates : {false, true}) {
auto prepared = prepareSpatialSchedules(*module, target, 1, collectCandidates);
TEST_CHECK(succeeded(prepared) && prepared->candidates.size() == 1);
TEST_CHECK(prepared->candidates.front().dominanceOrderCompute.empty());
}
}
static void testSelectedLoweringFailure() {
using namespace mlir;
MLIRContext context(MLIRContext::Threading::DISABLED);
context.loadDialect<SpatialDialect, pim::PimDialect, func::FuncDialect,
arith::ArithDialect, tensor::TensorDialect, memref::MemRefDialect,
scf::SCFDialect, linalg::LinalgDialect, affine::AffineDialect,
bufferization::BufferizationDialect>();
OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
OpBuilder builder(module->getBodyRegion());
auto loc = module->getLoc();
auto type = RankedTensorType::get({1, 536870912}, builder.getF32Type());
auto function = func::FuncOp::create(builder, loc, "main_graph", builder.getFunctionType({type}, {type}));
builder.setInsertionPointToStart(function.addEntryBlock());
auto compute = SpatGraphCompute::create(builder, loc, TypeRange{type}, ValueRange{},
ValueRange{function.getArgument(0)});
compute.getBody().push_back(new Block);
Value input = compute.getBody().front().addArgument(type, loc);
builder.setInsertionPointToStart(&compute.getBody().front());
Value result = SpatReluOp::create(builder, loc, type, input);
SpatYieldOp::create(builder, loc, ValueRange{result});
builder.setInsertionPointAfter(compute);
func::ReturnOp::create(builder, loc, compute.getResults());
TEST_CHECK(succeeded(verify(*module)));
SchedulingTarget target;
target.processorCount = 2;
target.residentWeightCapacity = 1;
target.synchronizationRegisterCount = 64;
target.interProcessorLatencyNs = {0, 2, 2, 0};
target.averageInterProcessorLatencyNs = 2;
target.estimateRealizedScheduleCost = [](ModuleOp, llvm::json::Object &) {
TEST_CHECK(false && "single-candidate compilation must skip costs");
return WholeScheduleCost{};
};
SpatialTargetResources resources{{128, 128}, 1, 2, 16};
std::string diagnostics;
ScopedDiagnosticHandler handler(&context, [&](Diagnostic &diagnostic) {
diagnostics += diagnostic.str(); return success();
});
ScopedCompilerOutputSuppression silence;
// Match the compiler's preloaded backend registry before nested pipelines run.
PassManager backend(&context);
addPimLoweringPasses(backend, resources, EmitPimCodegen, false);
DialectRegistry registry;
backend.getDependentDialects(registry);
pim::registerOpBufferizationInterfaces(registry);
arith::registerBufferizableOpInterfaceExternalModels(registry);
bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(registry);
context.appendDialectRegistry(registry);
context.loadAllAvailableDialects();
for (bool collectCandidates : {false, true}) {
OwningOpRef<ModuleOp> checkpoint = cast<ModuleOp>((*module)->clone());
auto prepared = prepareSpatialSchedules(*checkpoint, target, 2, collectCandidates);
TEST_CHECK(succeeded(prepared) && !prepared->candidates.empty());
}
PassManager manager(&context);
manager.addPass(createPimPipelineSelectionPass(target, resources, 2, SpatialDataflowExportStage::None));
TEST_CHECK(failed(manager.run(*module)));
TEST_CHECK(diagnostics.find("backend selection failed") != std::string::npos);
TEST_CHECK(diagnostics.find("outside representable range") != std::string::npos);
}
static void testResourceRetries() {
using namespace mlir;
auto previousCores = coresCount.getValue();
auto previousSize = crossbarSize.getValue();
auto previousCrossbars = crossbarCountInCore.getValue();
coresCount = 2; crossbarSize = 128; crossbarCountInCore = 1;
MLIRContext context;
context.loadDialect<pim::PimDialect, memref::MemRefDialect, func::FuncDialect,
arith::ArithDialect>();
OwningOpRef<ModuleOp> checkpoint = ModuleOp::create(UnknownLoc::get(&context));
OpBuilder builder(checkpoint->getBodyRegion());
auto loc = checkpoint->getLoc();
auto function = func::FuncOp::create(builder, loc, "main_graph", builder.getFunctionType({}, {}));
builder.setInsertionPointToStart(function.addEntryBlock());
auto core = pim::PimCoreOp::create(builder, loc, ValueRange{}, uint32_t{0});
core.getBody().push_back(new Block);
builder.setInsertionPointToStart(&core.getBody().front());
auto allocation = memref::AllocOp::create(builder, loc,
MemRefType::get({int64_t{1} << 31}, builder.getI8Type()));
auto smallType = MemRefType::get({1}, builder.getI8Type());
auto destination = memref::AllocOp::create(builder, loc, smallType);
Value zero = arith::ConstantIndexOp::create(builder, loc, 0);
pim::PimMemCopyOp::create(builder, loc, smallType, zero, zero, destination,
allocation, builder.getI32IntegerAttr(1));
pim::PimHaltOp::create(builder, loc);
builder.setInsertionPointAfter(core);
func::ReturnOp::create(builder, loc);
TEST_CHECK(succeeded(verify(*checkpoint)));
std::string before;
llvm::raw_string_ostream stream(before); checkpoint->print(stream, OpPrintingFlags().printGenericOpForm());
auto original = scheduleCost(100);
auto faster = scheduleCost(90);
std::vector<WholeScheduleCost> costs{original, faster};
std::string reason;
SpatialTargetResources resources{{128, 128}, 1, 2, 16};
PassManager backend(&context);
addPimLoweringPasses(backend, resources, EmitPimCodegen, false);
DialectRegistry registry;
backend.getDependentDialects(registry);
pim::registerOpBufferizationInterfaces(registry);
arith::registerBufferizableOpInterfaceExternalModels(registry);
bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(registry);
context.appendDialectRegistry(registry);
context.loadAllAvailableDialects();
auto emit = [&](ModuleOp module) {
llvm::SmallString<128> temporary;
TEST_CHECK(!llvm::sys::fs::createUniqueDirectory("raptor-retry-program", temporary));
std::string output = temporary.str().str();
ScopedCompilerOutputSuppression silence;
TEST_CHECK(compileToPimCode(module, output) == CompilerSuccess);
auto binary = llvm::MemoryBuffer::getFile(output + "/core_0.pim");
TEST_CHECK(binary && (*binary)->getBufferSize() > 12);
std::string program = (*binary)->getBuffer().str();
TEST_CHECK(!llvm::sys::fs::remove_directories(temporary));
return program;
};
std::string emitted;
for (bool allFail : {false, true}) {
std::vector<size_t> attempted;
OwningOpRef<ModuleOp> winner;
std::vector<llvm::json::Object> summaries;
auto selected = lowerSelectedPimCandidate(costs, [&](size_t index) {
attempted.push_back(index);
IRMapping mapping;
OwningOpRef<ModuleOp> candidate = cast<ModuleOp>((*checkpoint)->clone(mapping));
if (!index && !allFail)
cast<memref::AllocOp>(mapping.lookup(allocation.getResult()).getDefiningOp())
.getResult().setType(MemRefType::get({16}, builder.getI8Type()));
llvm::json::Object summary;
auto result = lowerPimPipelineCandidate(*candidate, resources, summary);
if (result == PimCandidateResult::Success) winner = std::move(candidate);
summaries.push_back(std::move(summary));
return result;
}, reason);
TEST_CHECK(attempted == std::vector<size_t>({1, 0}));
TEST_CHECK(allFail ? !selected : selected == 0);
TEST_CHECK(summaries.front().getString("status") == "resource-limit");
TEST_CHECK(summaries.front().getString("reason")->contains("signed int32 address range"));
if (selected) emitted = emit(*winner);
if (allFail) TEST_CHECK(reason.find("all candidate schedules failed") != std::string::npos);
std::string after;
llvm::raw_string_ostream out(after); checkpoint->print(out, OpPrintingFlags().printGenericOpForm());
TEST_CHECK(before == after);
}
TEST_CHECK(!emitted.empty());
IRMapping mapping;
OwningOpRef<ModuleOp> alternative = cast<ModuleOp>((*checkpoint)->clone(mapping));
cast<memref::AllocOp>(mapping.lookup(allocation.getResult()).getDefiningOp())
.getResult().setType(MemRefType::get({16}, builder.getI8Type()));
llvm::json::Object summary;
TEST_CHECK(lowerPimPipelineCandidate(*alternative, resources, summary) == PimCandidateResult::Success);
TEST_CHECK(emit(*alternative) == emitted);
size_t attempts = 0;
auto empty = lowerSelectedPimCandidate({}, [&](size_t) {
++attempts;
return PimCandidateResult::Success;
}, reason);
TEST_CHECK(!empty && attempts == 0 && reason.find("at least one candidate") != std::string::npos);
auto selected = lowerSelectedPimCandidate(costs, [&](size_t) {
++attempts;
OwningOpRef<ModuleOp> invalid = cast<ModuleOp>((*alternative)->clone());
invalid->walk([&](pim::PimMemCopyOp copy) {
copy->setAttr("size", builder.getI64IntegerAttr(1));
});
llvm::json::Object failure;
auto result = lowerPimPipelineCandidate(*invalid, resources, failure);
TEST_CHECK(result == PimCandidateResult::Failure);
TEST_CHECK(failure.getString("status") == "systemic-or-unclassified-error");
return result;
}, reason);
TEST_CHECK(!selected && attempts == 1);
TEST_CHECK(reason.find("recovery is not justified") != std::string::npos);
for (unsigned unknown = 0; unknown < 4; ++unknown) {
auto invalid = faster;
if (unknown == 0) invalid.primary.reset();
if (unknown == 1) invalid.recurrence[0].status = ScheduleRecurrenceEvidence::Unproved;
if (unknown == 2) invalid.loweringSupport = StaticLoweringSupport::Unknown;
if (unknown == 3) invalid.loweringSupport = StaticLoweringSupport::Unsupported;
attempts = 0;
std::vector<WholeScheduleCost> incomplete{original, invalid};
selected = lowerSelectedPimCandidate(incomplete, [&](size_t) {
++attempts;
return PimCandidateResult::Success;
}, reason);
TEST_CHECK(!selected && attempts == 0);
}
coresCount = previousCores; crossbarSize = previousSize; crossbarCountInCore = previousCrossbars;
}
int main() {
testEmptyGraphPreparation();
testSelectedLoweringFailure();
testResourceRetries();
mlir::MLIRContext context;
context.allowUnregisteredDialects();
mlir::OwningOpRef<mlir::ModuleOp> checkpoint = mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
@@ -25,56 +273,52 @@ int main() {
mlir::OwningOpRef<mlir::ModuleOp> trial = mlir::cast<mlir::ModuleOp>((*checkpoint)->clone(mapping));
auto remapped = spatial::remapSpatialSchedule(schedule, mapping);
auto copy = remapped.dominanceOrderCompute.front();
assert(copy.op != op && copy.laneStart == 3 && copy.laneCount == 2);
assert(remapped.computeToCpuMap.lookup(copy) == 7 && remapped.computeToCpuSlotMap.lookup(copy) == 4);
assert(remapped.computeToAestMap.lookup(copy) == 42 && remapped.isLastComputeOfCpu.contains(copy));
assert(remapped.cpuToLastComputeMap.lookup(7) == copy);
TEST_CHECK(copy.op != op && copy.laneStart == 3 && copy.laneCount == 2);
TEST_CHECK(remapped.computeToCpuMap.lookup(copy) == 7 && remapped.computeToCpuSlotMap.lookup(copy) == 4);
TEST_CHECK(remapped.computeToAestMap.lookup(copy) == 42 && remapped.isLastComputeOfCpu.contains(copy));
TEST_CHECK(remapped.cpuToLastComputeMap.lookup(7) == copy);
copy.op->erase();
assert(schedule.dominanceOrderCompute.front().op == op && op->getBlock() == checkpoint->getBody());
TEST_CHECK(schedule.dominanceOrderCompute.front().op == op && op->getBlock() == checkpoint->getBody());
}
auto score = [](double value, std::string id) { return WholeScheduleCost{value, value, {}, id, 0.021}; };
auto score = [](double value, std::string) {
return scheduleCost(value);
};
std::string reason;
auto original = score(100, "0"), faster = score(90, "1");
for (auto support : {StaticLoweringSupport::Unknown,
StaticLoweringSupport::Supported, StaticLoweringSupport::Unsupported}) {
auto cost = faster;
cost.loweringSupport = support;
TEST_CHECK(cost.available());
TEST_CHECK(cost.eligible() == (support == StaticLoweringSupport::Supported));
TEST_CHECK(cost.loweringSupported().has_value() == (support != StaticLoweringSupport::Unknown));
}
auto incompleteProof = faster;
incompleteProof.recurrence[1].periodGenerations = 0;
TEST_CHECK(incompleteProof.available() && !incompleteProof.eligible());
incompleteProof.recurrence[1] = {};
TEST_CHECK(incompleteProof.convergenceProved() == false && !incompleteProof.eligible());
incompleteProof.recurrence[0] = {};
TEST_CHECK(!incompleteProof.convergenceProved().has_value());
auto select = [&](std::initializer_list<WholeScheduleCost> values) {
return selectPipelineCandidate(llvm::ArrayRef<WholeScheduleCost>(values), reason);
};
assert(select({original, faster}) == 1);
assert(select({original, score(97, "2")}) == 0);
assert(select({original, score(105, "2")}) == 0);
TEST_CHECK(select({original, faster}) == 1);
TEST_CHECK(select({original, score(97, "2")}) == 0);
TEST_CHECK(select({original, score(105, "2")}) == 0);
auto sensitive = faster; sensitive.alternative = 101;
assert(select({original, sensitive}) == 0);
assert(select({original, faster, score(91, "2")}) == 1);
assert(select({original, faster, score(90, "1")}) == 1);
TEST_CHECK(select({original, sensitive}) == 0);
TEST_CHECK(select({original, faster, score(91, "2")}) == 1);
TEST_CHECK(select({original, faster, score(90, "1")}) == 1);
auto unavailable = WholeScheduleCost{};
assert(select({unavailable, faster}) == 0);
assert(select({original, unavailable}) == 0);
TEST_CHECK(select({unavailable, faster}) == 0);
TEST_CHECK(select({original, unavailable}) == 0);
auto unknownUncertainty = faster;
unknownUncertainty.relativeUncertainty.reset();
assert(select({original, unknownUncertainty}) == 0);
unknownUncertainty.selectionMargin.reset();
TEST_CHECK(select({original, unknownUncertainty}) == 0);
for (double value : {0., -1., std::numeric_limits<double>::infinity(), std::numeric_limits<double>::quiet_NaN()})
assert(select({original, score(value, "2")}) == 0);
for (int i = 0; i < 8; ++i) assert(select({original, faster}) == 1);
TEST_CHECK(select({original, score(value, "2")}) == 0);
for (int i = 0; i < 8; ++i) TEST_CHECK(select({original, faster}) == 1);
// The generic query accepts target-lowered IR without a target dialect API.
for (double interval : {90., 97.}) {
SchedulingTarget target;
assert(!target.prepareScheduleCost);
target.prepareScheduleCost = [interval](mlir::ModuleOp module,
llvm::json::Object &) -> mlir::FailureOr<ScheduleCostQuery> {
assert(module.getBody()->getOperations().size() == 1);
ScheduleCostQuery query;
query.identity = "static-program";
query.evaluate = [interval](llvm::json::Object &) {
return WholeScheduleCost{interval, interval, {}, "static-program", 0.021};
};
return query;
};
llvm::json::Object summary;
auto query = target.prepareScheduleCost(*checkpoint, summary);
assert(mlir::succeeded(query) && query->identity == "static-program");
auto cost = query->evaluate(summary);
assert(cost.eligible() && select({original, cost}) == (interval == 90. ? 1 : 0));
assert(checkpoint->getBody()->front().getName().getStringRef() == "test.compute");
}
}
test/PIM/ProjectionMemoryCostTest.cpp
+99
View File
@@ -0,0 +1,99 @@
#include "TestCheck.hpp"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/Diagnostics.h"
#include "mlir/IR/Verifier.h"
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.hpp"
#include <limits>
using namespace mlir;
using namespace onnx_mlir;
using namespace onnx_mlir::spatial;
static void checkProjection(ArrayRef<int64_t> shape, ArrayRef<int64_t> sizes,
ArrayRef<int64_t> strides, bool available, bool valid = true) {
MLIRContext context(MLIRContext::Threading::DISABLED);
context.loadDialect<SpatialDialect, tensor::TensorDialect, func::FuncDialect>();
OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
OpBuilder builder(module->getBodyRegion());
auto loc = module->getLoc();
auto sourceType = RankedTensorType::get(shape, builder.getF32Type());
auto resultType = RankedTensorType::get(sizes, builder.getF32Type());
auto function = func::FuncOp::create(builder, loc, "projection",
builder.getFunctionType({sourceType}, {resultType}));
builder.setInsertionPointToStart(function.addEntryBlock());
auto compute = createEmptySpatGraphCompute(builder, loc, TypeRange{resultType},
ValueRange{}, ValueRange{function.getArgument(0)});
auto attrs = [&](ArrayRef<int64_t> values) {
SmallVector<OpFoldResult> result;
for (int64_t value : values) result.push_back(builder.getIndexAttr(value));
return result;
};
auto slice = tensor::ExtractSliceOp::create(builder, loc, resultType,
*compute.getInputArgument(0), attrs(SmallVector<int64_t>(shape.size(), 0)),
attrs(sizes), attrs(strides));
Value result = SpatReluOp::create(builder, loc, resultType, slice);
SpatYieldOp::create(builder, loc, ValueRange{result});
builder.setInsertionPointAfter(compute);
func::ReturnOp::create(builder, loc, compute.getResults());
std::string diagnostics;
ScopedDiagnosticHandler handler(&context, [&](Diagnostic &diagnostic) {
diagnostics += diagnostic.str();
return success();
});
bool verified = succeeded(verify(*module));
if (verified != valid) std::fprintf(stderr, "%s\n", diagnostics.c_str());
TEST_CHECK(verified == valid);
if (!valid) return;
diagnostics.clear();
for (Cost width : {1, 16}) {
diagnostics.clear();
SchedulingTarget target;
target.processorCount = 1;
target.sharedMemoryWidthBytes = width;
target.sharedMemoryReadLatencyNs = width == 1 ? 7 : 3;
auto graph = buildComputeGraph(function, target);
TEST_CHECK(succeeded(graph));
if (available) {
TEST_CHECK(graph->baselineMemory.readRequests == 2);
TEST_CHECK(graph->baselineMemory.readBytes == 4);
TEST_CHECK(graph->baselineMemory.serviceNs == (width == 1 ? 28 : 6));
TEST_CHECK(diagnostics.empty());
} else {
TEST_CHECK(graph->baselineMemory.readRequests == static_cast<Cost>(resultType.getNumElements()));
TEST_CHECK(graph->baselineMemory.readBytes == static_cast<Cost>(resultType.getNumElements()));
TEST_CHECK(graph->baselineMemory.serviceNs == resultType.getNumElements()
* target.sharedMemoryReadLatencyNs);
auto message = diagnostics.find("projection memory cost unavailable");
TEST_CHECK(message != std::string::npos);
TEST_CHECK(diagnostics.find("projection memory cost unavailable", message + 1) == std::string::npos);
}
// The optional estimate must not make otherwise verified Spatial IR fail
// schedule preparation. Target lowering retains its own legality checks.
target.residentWeightCapacity = 1;
target.interProcessorLatencyNs = {0};
function.setName("main_graph");
auto schedules = prepareSpatialSchedules(*module, target, 1, false);
TEST_CHECK(succeeded(schedules) && schedules->candidates.size() == 1);
IRMapping mapping;
OwningOpRef<ModuleOp> scheduled = cast<ModuleOp>(module->getOperation()->clone(mapping));
auto schedule = remapSpatialSchedule(schedules->candidates.front(), mapping);
auto realized = realizeSpatialSchedule(*scheduled, schedule, target, 1,
SpatialDataflowExportStage::None);
if (failed(realized)) std::fprintf(stderr, "%s\n", diagnostics.c_str());
TEST_CHECK(succeeded(realized));
TEST_CHECK(succeeded(verify(*scheduled)));
}
}
int main() {
int64_t limit = std::numeric_limits<int64_t>::max();
checkProjection({3, 5}, {2, 2}, {1, 1}, true);
// These tensor types and slices are valid Spatial IR. The cost query cannot
// represent their flattened offsets or strides with signed 64-bit values.
checkProjection({2, limit}, {2, 2}, {1, 1}, false);
checkProjection({1, limit, 2}, {1, 1, 1}, {1, 1, 1}, false);
checkProjection({2, limit}, {1, 1}, {2, 1}, false);
checkProjection({3, 5}, {4, 2}, {1, 1}, false, false);
}
test/PIM/SpatialSchedulingTargetTest.cpp
+1
View File
@@ -1048,6 +1048,7 @@ static void checkPublicationMemoryTraffic(mlir::func::FuncOp function,
auto report = getPeriodicScheduleReport(graph, *materialized, *shared, *boundaries, timed);
assert(report.getString("status") == "estimated");
assert(report.getBoolean("diagnostic_only") == true);
assert(report.getBoolean("selection_eligible") == false);
assert(report.getInteger("boundary_write_bytes") == expected.writeBytes);
assert(report.getInteger("boundary_read_bytes") == expected.readBytes);
timed.sharedMemoryReadReturnsPayload = false;
test/PIM/TestCheck.hpp
+11
View File
@@ -0,0 +1,11 @@
#pragma once
#include <cstdio>
#include <cstdlib>
#define TEST_CHECK(condition) do { \
if (!(condition)) { \
std::fprintf(stderr, "%s:%d: check failed: %s\n", __FILE__, __LINE__, #condition); \
std::exit(EXIT_FAILURE); \
} \
} while (false)
test/PIM/test_PimEmission.py
+5 -3
View File
@@ -53,9 +53,7 @@ def main():
report = json.loads((out / "reports/pipeline_selection.json").read_text())
assert report["generated_candidates"] == 1
assert report["single_candidate_analysis_skipped"]
assert report["analyzed_candidates"] == report["lowered_candidates"] == 0
assert report["unavailable_estimates"] == 0
assert not report["reused_lowered_candidate"]
assert report["lowered_candidates"] == 1
assert "no cost comparison required" in report["reason"]
else:
assert "intermediate emission retains the original schedule" in output
@@ -83,6 +81,10 @@ def main():
capture_output=True, text=True, timeout=30)
assert result.returncode == 0, result.stdout + result.stderr
assert (out / "reports/pipeline_periodic.json").exists()
report = json.loads((out / "reports/pipeline_periodic.json").read_text())
assert report["diagnostic_only"]
assert not report["selection_eligible"]
assert not (out / "pim").exists()
if __name__ == "__main__":
test/PIM/test_PimPeriodicDiagnostic.py
+4 -1
View File
@@ -69,8 +69,11 @@ def check_controller(run):
controller=True, issue=True)
reader = next(core for core in zero['memory_controller']['per_core'] if core['core'] == 0)
assert reader['queue_wait_ns'] > 0 # Zero service still queues behind an active request.
bounded = run([[scalar, scalar], [scalar, load(65536), scalar]], controller=True, issue=True)
bounded = json.loads(run([[scalar, scalar], [scalar, load(65536), scalar]],
controller=True, issue=True, valid=False).stdout)
assert bounded['status'] == 'unavailable'
assert bounded['memory_controller']['status'] == 'unavailable'
assert 'periodic_state' not in bounded['memory_controller']
shared = [[record(31, a=0, b=2), store(128), record(32, arg=1, a=1), record(32, arg=2, a=1)],
[record(32, arg=0, a=0), record(31, a=1, b=1), load(128), record(29, arg=2, length=64)],
[record(32, arg=0, a=0), record(31, a=1, b=1), load(128), record(30, arg=1, length=64)]]
test/PIM/test_PimSpatialSelection.py
+112
View File
@@ -0,0 +1,112 @@
"""Spatial selection lowers one winner and rejects incomplete comparisons."""
import argparse
import hashlib
import json
import os
from pathlib import Path
import shutil
import subprocess
import tempfile
import onnx
from onnx import TensorProto, helper
REPO = Path(__file__).resolve().parents[2]
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--compiler", required=True)
args = parser.parse_args()
with tempfile.TemporaryDirectory(prefix="raptor-spatial-selection-") as temporary:
root = Path(temporary)
guard = root / "static-audit.so"
subprocess.run([os.environ.get("CC", "cc"), "-shared", "-fPIC",
str(REPO / "test/PIM/StaticSchedulingAudit.c"), "-o", str(guard)], check=True)
env = dict(os.environ, LD_PRELOAD=str(guard), OMP_NUM_THREADS="1")
iteration = 0
def compile(source, pipeline=2, emit="EmitPimCodegen", extra=(), succeeds=True):
nonlocal iteration
output = root / str(iteration)
iteration += 1
output.mkdir()
command = [args.compiler, str(source), "-o", str(output / "model"), "--maccel=PIM",
f"--{emit}", f"--pipeline={pipeline}", "--core-count=138", "--crossbar-count=128",
"--crossbar-size=128", "--mlir-disable-threading", "--pim-detect-communication-deadlock",
f"--pim-target-config={REPO}/validation/pimsim_configs/pimcomp/arch-b/throughput_config_1000ms.json"]
command.extend(extra)
result = subprocess.run(command, env=env, capture_output=True, text=True, timeout=120)
assert (result.returncode == 0) == succeeds, result.stdout + result.stderr
path = output / "reports/pipeline_selection.json"
report = json.loads(path.read_text()) if path.exists() else None
program = {str(p.relative_to(output / "pim")): hashlib.sha256(p.read_bytes()).hexdigest()
for p in (output / "pim").rglob("*") if p.is_file()}
if "--pim-export-spatial-dataflow=all" in extra:
for stage in ("spatial3_scheduled_no_comm", "spatial4_scheduled"):
for suffix in ("nodes.csv", "edges.csv"):
assert (output / f"reports/{stage}.{suffix}").stat().st_size > 0
assert (output / "dialects/spatial4_scheduled.mlir").exists()
assert (output / "dialects/pim1_buff.mlir").exists()
shutil.rmtree(output)
return report, program, result.stdout + result.stderr
vgg = REPO / "validation/networks/pimcomp_models/vgg8/vgg8-mnist-reconstructed.onnx"
selected, program, _ = compile(vgg)
count = selected["generated_candidates"]
assert count > 1 and selected["strategy"] == "spatial-queue-aware"
assert selected["lowered_candidates"] == 1
assert len(selected["backend_attempts"]) == 1
assert selected["backend_attempts"][0]["candidate"] == selected["selection"]
assert not selected["prediction_error_bound_available"]
assert selected["selection_margin"] == 0.021
for candidate in selected["candidates"]:
assert candidate["cost_available"] and candidate["convergence_proved"]
assert candidate["lowering_supported"] and candidate["eligible"]
assert "request_order" not in candidate["spatial_cost"]
assert "without_local_copy_work" not in candidate["spatial_cost"]
diagnostic, diagnostic_program, _ = compile(vgg, extra=(
"--pim-export-spatial-dataflow=all", "--pim-dump-ir", "--pim-periodic-diagnostics"))
assert diagnostic_program == program and diagnostic["selection"] == selected["selection"]
# A single schedule needs no optional timing contract.
source = root / "softmax.onnx"
tensor = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 8])
output = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 8])
model = helper.make_model(helper.make_graph([helper.make_node("Softmax", ["x"], ["y"], axis=1)],
"softmax", [tensor], [output]), opset_imports=[helper.make_opsetid("", 18)])
model.ir_version = 10
onnx.save(model, source)
selected, program, _ = compile(source)
assert selected["single_candidate_analysis_skipped"] and selected["lowered_candidates"] == 1
assert "spatial_cost" not in selected["candidates"][0]
# An unknown estimate cannot silently exclude any alternative.
model = onnx.load(vgg)
output = model.graph.output[0].name
internal = output + "_before_softmax"
for node in model.graph.node:
for index, value in enumerate(node.output):
if value == output:
node.output[index] = internal
for index, value in enumerate(node.input):
if value == output:
node.input[index] = internal
model.graph.node.append(helper.make_node("Softmax", [internal], [output], axis=-1))
source = root / "vgg-softmax.onnx"
onnx.save(model, source)
report, program, diagnostics = compile(source, succeeds=False)
assert report["generated_candidates"] > 1 and report["status"] == "incomplete-costs"
assert "cannot compare the complete candidate set" in diagnostics
assert "unsupported abstract operation" in diagnostics
assert not program and "backend_attempts" not in report
# Latency PEFT and earlier emission do not invoke candidate costing.
report, program, _ = compile(vgg, pipeline=1)
assert report is None and program
report, program, stderr = compile(vgg, emit="EmitSpatial")
assert report is None and not program
assert "intermediate emission retains the original schedule" in stderr
if __name__ == "__main__":
main()
test/PIM/test_QueueSelection.py
+14 -21
View File
@@ -13,16 +13,7 @@ REPO = Path(__file__).resolve().parents[2]
def check_selected_program(out):
report = json.loads((out / "reports/pipeline_selection.json").read_text())
program = json.loads((out / "pim/config.json").read_text())
digest = hashlib.sha256()
for core in range(program["core_cnt"]):
digest.update((out / f"pim/core_{core}.pim").read_bytes())
for count in program["array_group_map"][f"core{core}"]:
digest.update(f"{count},".encode())
digest.update(b";")
assert digest.hexdigest().upper() == report["candidates"][report["selection"]]["identity"], \
"emitted program differs from the analyzed winner"
assert report["reused_lowered_candidate"]
assert report["strategy"] == "spatial-queue-aware" and report["lowered_candidates"] == 1
for name in ("spatial3_scheduled_no_comm.txt", "memory_report.txt"):
assert (out / "reports" / name).stat().st_size > 0, name
return report
@@ -71,12 +62,13 @@ def main():
assert result.returncode == 0, result.stdout + result.stderr
report = json.loads((out / "reports/pipeline_selection.json").read_text())
assert report["static_analysis"] and report["temporary_disk_bytes"] == 0
assert report["analyzed_candidates"] > 1, report
assert report["generated_candidates"] > 1, report
assert all("identity" not in c and "request_order" not in c["spatial_cost"]
for c in report["candidates"])
check_selected_program(out)
assert not report["reconstructed_spatial_artifacts"]
assert not (out / "dialects").exists()
assert not list(out.glob("candidate*"))
decisions.append((report["selection"], [(c["identity"], c.get("interval_ns"), c.get("alternative_tie_interval_ns"))
decisions.append((report["selection"], [(c.get("interval_ns"), c.get("alternative_tie_interval_ns"))
for c in report["candidates"]]))
# Keep the test's disk peak to one selected program.
shutil.rmtree(out)
@@ -95,8 +87,9 @@ def main():
result = subprocess.run(command, cwd=root, env=env, capture_output=True, text=True, timeout=90)
assert result.returncode == 0, result.stdout + result.stderr
report = check_selected_program(out)
assert report["selection"] != 0 and report["reconstructed_spatial_artifacts"] == diagnostics
identity = report["candidates"][report["selection"]]["identity"]
assert report["selection"] != 0 and bool((out / "dialects").exists()) == diagnostics
identity = {str(p.relative_to(out / "pim")): hashlib.sha256(p.read_bytes()).hexdigest()
for p in (out / "pim").rglob("*") if p.is_file()}
schedule = (out / "reports/spatial3_scheduled_no_comm.txt").read_bytes()
if diagnostics:
assert identity == previous_identity and schedule == previous_schedule
@@ -123,9 +116,9 @@ def main():
result = subprocess.run(command, cwd=root, env=env, capture_output=True, text=True, timeout=60)
assert result.returncode == 0, result.stdout + result.stderr
report = json.loads((out / "reports/pipeline_selection.json").read_text())
assert report["selection"] == 0 and report["analyzed_candidates"] == 0
assert report["generated_candidates"] == 1 and report["lowered_candidates"] == 0
assert "ROB capacity one" in report["reason"] and report["unavailable_estimates"] > 0
assert report["selection"] == 0 and report["single_candidate_analysis_skipped"]
assert report["generated_candidates"] == 1 and report["lowered_candidates"] == 1
assert "no cost comparison required" in report["reason"]
assert report["temporary_disk_bytes"] == 0 and not list(root.glob("candidate*"))
shutil.rmtree(out)
out.mkdir()
@@ -134,9 +127,9 @@ def main():
result = subprocess.run(command, cwd=root, env=env, capture_output=True, text=True, timeout=60)
assert result.returncode == 0, result.stdout + result.stderr
report = json.loads((out / "reports/pipeline_selection.json").read_text())
assert report["selection"] == 0 and report["analyzed_candidates"] == 0
assert report["generated_candidates"] == 1 and report["lowered_candidates"] == 0
assert "synchronization ablation" in report["reason"]
assert report["selection"] == 0 and report["single_candidate_analysis_skipped"]
assert report["generated_candidates"] == 1 and report["lowered_candidates"] == 1
assert "no cost comparison required" in report["reason"]
assert not (root / "executed").exists(), "production compiler executed a diagnostic or simulator"
test/PIM/test_ShapeUtils.py
+11
View File
@@ -0,0 +1,11 @@
"""Overflow must fail explicitly, including builds with disabled assertions."""
import argparse
import subprocess
parser = argparse.ArgumentParser()
parser.add_argument("--test", required=True)
args = parser.parse_args()
result = subprocess.run([args.test, "--stride-overflow"], capture_output=True, text=True)
assert result.returncode != 0
assert "row-major strides are not representable in int64_t" in result.stderr, result.stderr
assert "Assertion" not in result.stderr