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

2026-10-26 12:07:01 +01:00
parent fd493ae6e8
commit 4b53f6cf10
62 changed files with 4293 additions and 871 deletions
@@ -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);
+}