NiccoloN
a77edafcbf
Validate Operations / validate-operations (push) Has been cancelled
plot adjustments validation adjustments
100 lines
4.9 KiB
C++
100 lines
4.9 KiB
C++
#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);
|
|
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);
|
|
}
|