better scheduling

2026-10-13 11:53:03 +02:00
parent 80bbf75883
commit 210c9a8363
6 changed files with 443 additions and 62 deletions
@@ -4,14 +4,165 @@
 #include <vector>

 #include "mlir/IR/BuiltinOps.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"

 #include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredTransferPlanning.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.hpp"

 using namespace onnx_mlir::spatial;

+static void testResidentWeightViewCost() {
+  mlir::MLIRContext context;
+  context.loadDialect<SpatialDialect, mlir::tensor::TensorDialect>();
+  mlir::Block operands;
+  mlir::OwningOpRef<mlir::ModuleOp> owner =
+    mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
+  mlir::OpBuilder builder(&context);
+  auto loc = owner->getLoc();
+  auto weightType = mlir::RankedTensorType::get({128, 128}, builder.getF32Type());
+  auto inputType = mlir::RankedTensorType::get({1, 128}, builder.getF32Type());
+  auto weight = operands.addArgument(weightType, loc);
+  auto input = operands.addArgument(inputType, loc);
+  auto wideType = mlir::RankedTensorType::get({1, 256}, builder.getF32Type());
+  auto wideInput = operands.addArgument(wideType, loc);
+  auto makeCompute = [&](bool weightView, unsigned inputView) {
+    builder.setInsertionPointToEnd(owner->getBody());
+    auto compute = SpatGraphCompute::create(builder, loc,
+      mlir::TypeRange{inputType}, mlir::ValueRange{weight},
+      mlir::ValueRange{inputView == 2 ? wideInput : input});
+    auto *block = new mlir::Block();
+    compute.getBody().push_back(block);
+    mlir::Value w = block->addArgument(weightType, loc);
+    mlir::Value x = block->addArgument(inputView == 2 ? wideType : inputType, loc);
+    builder.setInsertionPointToStart(block);
+    auto slice = [&](mlir::Value value, mlir::RankedTensorType type, bool strided = false) {
+      llvm::SmallVector<mlir::OpFoldResult> offsets(2, builder.getIndexAttr(0));
+      llvm::SmallVector<mlir::OpFoldResult> strides(2, builder.getIndexAttr(1));
+      if (strided)
+        strides.back() = builder.getIndexAttr(2);
+      llvm::SmallVector<mlir::OpFoldResult> sizes;
+      for (int64_t size : type.getShape())
+        sizes.push_back(builder.getIndexAttr(size));
+      return mlir::tensor::ExtractSliceOp::create(
+        builder, loc, type, value, offsets, sizes, strides).getResult();
+    };
+    if (weightView)
+      w = slice(w, weightType);
+    if (inputView)
+      x = slice(x, inputType, inputView == 2);
+    auto result = SpatVMMOp::create(builder, loc, inputType, w, x);
+    SpatYieldOp::create(builder, loc, mlir::ValueRange{result.getResult()});
+    return ComputeInstance{compute, 0, 1};
+  };
+  auto direct = makeCompute(false, false);
+  auto weightView = makeCompute(true, false);
+  auto inputView = makeCompute(false, true);
+  auto stridedInput = makeCompute(false, 2);
+  for (Cost readLatency : {1, 20}) {
+    SchedulingTarget target;
+    target.localMemoryReadLatencyCycles = readLatency;
+    target.localMemoryWriteLatencyCycles = 3;
+    target.localMemoryWidthBytes = 8;
+    target.processorPeriodNs = 1;
+    target.computeBitwidth = 8;
+    assert(getComputeInstanceCost(weightView, target) == getComputeInstanceCost(direct, target));
+    assert(getComputeInstanceCost(inputView, target) == getComputeInstanceCost(direct, target));
+    assert(getComputeInstanceCost(stridedInput, target)
+      == getComputeInstanceCost(direct, target) + 64 * readLatency);
+  }
+}
+
+static void testPipelineStageSpan() {
+  using namespace mlir;
+  MLIRContext context;
+  context.loadDialect<SpatialDialect>();
+  OwningOpRef<ModuleOp> owner = ModuleOp::create(UnknownLoc::get(&context));
+  OpBuilder builder(owner->getBodyRegion());
+  ComputeGraph graph;
+  graph.successors.resize(12);
+  graph.predecessors.resize(12);
+  for (size_t group = 0; group < 6; ++group) {
+    auto batch = SpatGraphComputeBatch::create(builder, owner->getLoc(),
+      TypeRange{}, builder.getI32IntegerAttr(2), ValueRange{}, ValueRange{});
+    for (uint32_t lane = 0; lane < 2; ++lane) {
+      size_t task = graph.nodes.size();
+      ComputeInstance instance{batch, lane, 1};
+      graph.nodes.push_back({instance, group == 0 ? 40ULL : 10ULL, {}, task});
+      graph.instanceToIndex[instance] = task;
+      if (group != 0)
+        for (size_t predecessor = task - lane - 2; predecessor < task - lane; ++predecessor) {
+          graph.predecessors[task].push_back({predecessor, {}});
+          graph.successors[predecessor].push_back({task, {}});
+        }
+    }
+  }
+  for (size_t coreCount : {6, 12, 20}) {
+    SchedulingTarget target;
+    target.processorCount = coreCount;
+    target.residentWeightCapacity = 1;
+    target.interProcessorLatencyNs.assign(coreCount * coreCount, 0);
+    auto schedule = runPeftScheduler(graph, {target, &context});
+    std::string error;
+    assert(succeeded(applyPipelineScheduling(graph, schedule, 2, target, error)));
+    Cost stageCosts[2] = {};
+    Cost stageWork[2] = {};
+    for (size_t task = 0; task < graph.nodes.size(); ++task) {
+      size_t core = schedule.computeToCpuMap.lookup(graph.nodes[task].instance);
+      size_t stage = schedule.processorStages[core];
+      stageWork[stage] += graph.nodes[task].cost;
+      if (task % 2 == 0)
+        stageCosts[stage] += graph.nodes[task].cost;
+    }
+    // Both 40/50 and 50/40 are optimal cuts for this 90-cycle chain.
+    assert(std::max(stageCosts[0], stageCosts[1]) == 50);
+    if (coreCount == 20)
+      for (size_t stage = 0; stage < 2; ++stage)
+        assert(stageWork[stage] <= 10 * static_cast<size_t>(
+          llvm::count(schedule.processorStages, stage)));
+  }
+}
+
+static void testPipelinePackingCapacity() {
+  using namespace mlir;
+  MLIRContext context;
+  context.loadDialect<SpatialDialect>();
+  OwningOpRef<ModuleOp> owner = ModuleOp::create(UnknownLoc::get(&context));
+  OpBuilder builder(owner->getBodyRegion());
+  auto batch = SpatGraphComputeBatch::create(builder, owner->getLoc(),
+    TypeRange{}, builder.getI32IntegerAttr(4), ValueRange{}, ValueRange{});
+  ComputeGraph graph;
+  graph.successors.resize(4);
+  graph.predecessors.resize(4);
+  for (uint32_t task = 0; task < 4; ++task) {
+    ResidentWeightSet weights;
+    for (uint32_t index = 0; index < (task == 3 ? 3u : 1u); ++index) {
+      ResidentWeight weight;
+      weight.opaqueLane = task == 3 ? index + 1 : 0;
+      weights.push_back(weight);
+    }
+    ComputeInstance instance{batch, task, 1};
+    graph.nodes.push_back({instance, 10, weights, task});
+    graph.instanceToIndex[instance] = task;
+  }
+  SchedulingTarget target;
+  target.processorCount = 4;
+  target.residentWeightCapacity = 3;
+  target.interProcessorLatencyNs.assign(16, 0);
+  auto schedule = runPeftScheduler(graph, {target, &context});
+  std::string error;
+  assert(succeeded(applyPipelineScheduling(graph, schedule, 2, target, error)));
+  std::vector<Cost> loads(4);
+  for (const auto &node : graph.nodes) {
+    size_t core = schedule.computeToCpuMap.lookup(node.instance);
+    assert(schedule.processorStages[core] == 0);
+    loads[core] += node.cost;
+  }
+  assert(*std::max_element(loads.begin(), loads.end()) == 20);
+}
+
 static void testPhysicalPeft() {
  TransferCost transfer {.fixed = 50, .networkFlits = 4};
  SchedulingTarget line;
@@ -51,7 +202,7 @@ static void testPhysicalPeft() {
      chain.predecessors[task].push_back({task - 1, transfer});
    }
  }
-  line.residentWeightCapacity = reversed.residentWeightCapacity = 3;
+  line.residentWeightCapacity = reversed.residentWeightCapacity = 1;
  for (const SchedulingTarget &target : {line, reversed}) {
    auto scheduled = runPeftScheduler(chain, {target, &context});
    size_t firstCore = scheduled.computeToCpuMap.lookup(chain.nodes[0].instance);
@@ -69,6 +220,58 @@ static void testPhysicalPeft() {
  }
 }

+static void testPeftReservationCapacity() {
+  mlir::MLIRContext context;
+  auto owner = mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
+  ComputeGraph graph;
+  graph.successors.resize(4);
+  graph.predecessors.resize(4);
+  for (uint32_t task = 0; task < 4; ++task) {
+    ResidentWeightSet weights;
+    for (uint32_t index = 0; index < (task == 3 ? 4u : 1u); ++index) {
+      ResidentWeight weight;
+      weight.opaqueLane = task == 3 ? index + 1 : 0;
+      weights.push_back(weight);
+    }
+    graph.nodes.push_back({{owner, task, 1}, 10, weights, task});
+  }
+  SchedulingTarget target;
+  target.processorCount = 3;
+  target.residentWeightCapacity = 4;
+  target.interProcessorLatencyNs.assign(9, 0);
+  auto schedule = runPeftScheduler(graph, {target, &context});
+  size_t largeCore = schedule.computeToCpuMap.lookup(graph.nodes[3].instance);
+  for (size_t task : {0, 1, 2})
+    assert(schedule.computeToCpuMap.lookup(graph.nodes[task].instance) != largeCore);
+  owner.erase();
+}
+
+static void testPeftGapOrder() {
+  mlir::MLIRContext context;
+  mlir::OwningOpRef<mlir::ModuleOp> owner =
+    mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
+  ComputeGraph graph;
+  graph.successors.resize(3);
+  graph.predecessors.resize(3);
+  for (uint32_t task = 0; task < 3; ++task) {
+    ResidentWeight weight;
+    weight.opaqueLane = task == 0 ? 0 : 1;
+    graph.nodes.push_back({{owner->getOperation(), task, 1},
+      task == 2 ? 20u : 10u, {weight}, task});
+  }
+  TransferCost transfer {.fixed = 100};
+  graph.successors[0].push_back({1, transfer});
+  graph.predecessors[1].push_back({0, transfer});
+  SchedulingTarget target;
+  target.processorCount = 2;
+  target.residentWeightCapacity = 1;
+  target.interProcessorLatencyNs.assign(4, 0);
+  auto schedule = runPeftScheduler(graph, {target, &context});
+  assert(schedule.computeToCpuSlotMap.lookup(graph.nodes[2].instance) == 0);
+  assert(schedule.computeToCpuSlotMap.lookup(graph.nodes[1].instance) == 1);
+  assert(schedule.computeToAestMap.lookup(graph.nodes[2].instance) == 0);
+}
+
 static void testPipelineInputArrival() {
  ComputeGraph graph;
  graph.successors.resize(4);
@@ -146,9 +349,39 @@ static void testReadyCommunication() {
  assert(scheduled->slices.size() == 2);
  for (const auto &slice : scheduled->slices)
    assert(slice.sourceInsertionStep == 1);
+
+  plan.stepCounts = {3, 3, 3};
+  for (auto &exchange : plan.exchanges)
+    exchange->consumerStep = 2;
+  scheduled = scheduleDeferredCommunication({}, plan);
+  assert(mlir::succeeded(scheduled));
+  for (const auto &slice : scheduled->slices)
+    assert(slice.targetInsertionStep == 2);
+}
+
+static void testPipelineRestartOrder() {
+  mlir::MLIRContext context;
+  auto owner = mlir::ModuleOp::create(mlir::UnknownLoc::get(&context));
+  DeferredTransferPlan transfers;
+  for (int64_t core : {0, 7, 3, 5}) {
+    ScheduledInfo info;
+    info.op = owner;
+    info.cores.push_back(core);
+    info.pipelineStages.push_back(core == 0 ? 0 : 1);
+    transfers.scheduled.push_back(info);
+  }
+  auto boundaries = buildDeferredBoundaryPlan(transfers, {}, 16);
+  assert(mlir::succeeded(boundaries));
+  // Restarting the producer must not wait for a later consumer on core 3.
+  assert((transfers.downstreamCores == llvm::SmallVector<int64_t>{7, 3, 5}));
+  owner.erase();
 }

 int main() {
+  testResidentWeightViewCost();
+  testPipelineStageSpan();
+  testPipelinePackingCapacity();
+  testPipelineRestartOrder();
  testPipelineInputArrival();
  testReadyCommunication();
  PipelineCoreLayout unevenLayout(138, 4);
@@ -198,6 +431,8 @@ int main() {
  assert(slow.getInterProcessorLatencyNs(0, 1) == 10);

  testPhysicalPeft();
+  testPeftGapOrder();
+  testPeftReservationCapacity();

  ComputeGraph graph;
  graph.successors.resize(6);