automatic code reformat

src/PIM/Compiler/PimCodeGen.cpp

@@ -179,8 +179,7 @@ void PimMemory::allocateMemoryForValue(const MemoryValueKey& key, MemEntry& memE
   bool endFits = succeeded(checkedEnd) && *checkedEnd <= kPimAddressLimit;
   bool alignedEndFits = succeeded(checkedAlignedEnd) && *checkedAlignedEnd <= kPimAddressLimit;
   if (!startFits || !endFits || !alignedEndFits) {
-    printMemoryOverflowDiagnostic(
+    printMemoryOverflowDiagnostic(key.value,
-      key.value,
                                   key,
                                   memEntry.size,
                                   firstAvailableAddress,
@@ -218,8 +217,11 @@ PhysicalSlotInfo PimMemory::allocatePhysicalSlot(size_t slotSize, const MemoryVa
   bool endFits = succeeded(checkedEnd) && *checkedEnd <= kPimAddressLimit;
   bool alignedEndFits = succeeded(checkedAlignedEnd) && *checkedAlignedEnd <= kPimAddressLimit;
   if (!startFits || !endFits || !alignedEndFits) {
-    printMemoryOverflowDiagnostic(
+    printMemoryOverflowDiagnostic(key.value,
-      key.value, key, slot.size, firstAvailableAddress, succeeded(checkedAlignedEnd) ? *checkedAlignedEnd : kPimAddressLimit);
+                                  key,
+                                  slot.size,
+                                  firstAvailableAddress,
+                                  succeeded(checkedAlignedEnd) ? *checkedAlignedEnd : kPimAddressLimit);
     llvm_unreachable("PIM local memory allocation overflow");
   }
 
@@ -293,8 +295,8 @@ void PimMemory::allocateCore(Operation* op, std::optional<unsigned> lane) {
       const PhysicalSlotInfo& existingSlot = localPhysicalSlots[existingIndex];
       if (existingSlot.size < slot.requiredSize)
         continue;
-      auto candidateKey = std::tuple<size_t, size_t, size_t>(
+      auto candidateKey =
-        existingSlot.size - slot.requiredSize, existingSlot.size, existingSlot.id);
+        std::tuple<size_t, size_t, size_t>(existingSlot.size - slot.requiredSize, existingSlot.size, existingSlot.id);
       if (candidateKey < bestKey) {
         bestKey = candidateKey;
         bestExistingIndex = existingIndex;

src/PIM/Compiler/PimCompilerOptions.cpp

@@ -26,7 +26,8 @@ llvm::cl::opt<PimMemoryReportLevel> pimMemoryReport(
   "pim-memory-report",
   llvm::cl::desc("Emit a human-readable PIM memory planning report"),
   llvm::cl::values(clEnumValN(PimMemoryReportNone, "none", "Do not emit any PIM memory planning report")),
-  llvm::cl::values(clEnumValN(PimMemoryReportSummary, "summary", "Emit a concise slot reuse report with key offenders")),
+  llvm::cl::values(
+    clEnumValN(PimMemoryReportSummary, "summary", "Emit a concise slot reuse report with key offenders")),
   llvm::cl::values(clEnumValN(PimMemoryReportFull, "full", "Emit the full detailed PIM memory planning report")),
   llvm::cl::init(PimMemoryReportNone),
   llvm::cl::cat(OnnxMlirOptions));

src/PIM/Compiler/PimMemoryLiveness.cpp

@@ -5,8 +5,8 @@
 #include "mlir/Interfaces/DestinationStyleOpInterface.h"
 
 #include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/raw_ostream.h"
 
 #include <numeric>
@@ -222,8 +222,7 @@ struct OrderedTouchRange {
 
 static OrderedTouchRange
 getEffectiveTouchRange(mlir::Value definingValue, Operation* user, const OperationOrdering& ordering) {
-  OrderedTouchRange range {
+  OrderedTouchRange range {ordering.position.lookup(user), ordering.position.lookup(user), user, user, false};
-    ordering.position.lookup(user), ordering.position.lookup(user), user, user, false};
   for (Operation* current = user; current; current = current->getParentOp()) {
     auto forOp = dyn_cast<scf::ForOp>(current);
     if (!forOp || isWithin(definingValue, &forOp.getRegion()))
@@ -456,8 +455,7 @@ SmallVector<PlannedPhysicalSlot, 0> onnx_mlir::planPhysicalSlots(MutableArrayRef
   for (size_t intervalIndex : intervalOrder) {
     LocalAllocInterval& interval = intervals[intervalIndex];
     PlannedPhysicalSlot* bestSlot = nullptr;
-    auto bestKey = std::tuple<size_t, size_t, size_t, size_t>(
+    auto bestKey = std::tuple<size_t, size_t, size_t, size_t>(std::numeric_limits<size_t>::max(),
-      std::numeric_limits<size_t>::max(),
                                                               std::numeric_limits<size_t>::max(),
                                                               std::numeric_limits<size_t>::max(),
                                                               std::numeric_limits<size_t>::max());
@@ -476,8 +474,8 @@ SmallVector<PlannedPhysicalSlot, 0> onnx_mlir::planPhysicalSlots(MutableArrayRef
 
       size_t resultingSize = std::max(slot.requiredSize, interval.size);
       size_t growth = resultingSize - slot.requiredSize;
-      auto candidateKey = std::tuple<size_t, size_t, size_t, size_t>(
+      auto candidateKey =
-        growth, resultingSize, slot.intervalIndices.size(), slot.id);
+        std::tuple<size_t, size_t, size_t, size_t>(growth, resultingSize, slot.intervalIndices.size(), slot.id);
       if (candidateKey < bestKey) {
         bestKey = candidateKey;
         bestSlot = &slot;
@@ -553,7 +551,8 @@ MemoryPlanArtifacts onnx_mlir::buildMemoryPlanArtifacts(Operation *coreLikeOp,
     os << "Lane: " << *lane << "\n";
   os << "Summary:\n";
   os << "  logical allocation bytes:  " << formatReportMemory(totalLogicalBytes) << " (" << totalLogicalBytes << ")\n";
-  os << "  physical allocation bytes: " << formatReportMemory(totalPhysicalBytes) << " (" << totalPhysicalBytes << ")\n";
+  os << "  physical allocation bytes: " << formatReportMemory(totalPhysicalBytes) << " (" << totalPhysicalBytes
+     << ")\n";
   os << "  saved bytes:               " << formatReportMemory(savedBytes) << " (" << savedBytes << ")\n";
   os << "  saved percent: " << format("%.2f%%", savedPercent) << "\n";
   os << "  intervals: " << intervals.size() << "\n";
@@ -566,7 +565,8 @@ MemoryPlanArtifacts onnx_mlir::buildMemoryPlanArtifacts(Operation *coreLikeOp,
   os << "  largest logical allocation: " << largestLogicalAllocation << "\n";
   os << "  largest physical slot: " << largestPhysicalSlot << "\n";
   os << "  address limit: " << addressLimit << "\n";
-  os << "  peak physical memory: " << formatReportMemory(maximumAssignedAddress) << " (" << maximumAssignedAddress << ")\n";
+  os << "  peak physical memory: " << formatReportMemory(maximumAssignedAddress) << " (" << maximumAssignedAddress
+     << ")\n";
   os << "  maximum assigned address: " << maximumAssignedAddress << "\n";
 
   os << "\nHow To Read:\n";
@@ -577,12 +577,11 @@ MemoryPlanArtifacts onnx_mlir::buildMemoryPlanArtifacts(Operation *coreLikeOp,
 
   SmallVector<const PlannedPhysicalSlot*> reusedSlots;
   SmallVector<const PlannedPhysicalSlot*> singleUseSlots;
-  for (const PlannedPhysicalSlot &slot : slots) {
+  for (const PlannedPhysicalSlot& slot : slots)
     if (slot.intervalIndices.size() > 1)
       reusedSlots.push_back(&slot);
     else
       singleUseSlots.push_back(&slot);
-  }
 
   llvm::stable_sort(reusedSlots, [&](const PlannedPhysicalSlot* lhs, const PlannedPhysicalSlot* rhs) {
     uint64_t lhsLogicalBytes = getSlotLogicalBytes(*lhs, intervals);
@@ -607,18 +606,16 @@ MemoryPlanArtifacts onnx_mlir::buildMemoryPlanArtifacts(Operation *coreLikeOp,
   os << "\nBest Reuse:\n";
   if (reusedSlots.empty()) {
     os << "  no slots were shared by multiple intervals\n";
-  } else {
+  }
+  else {
     for (const PlannedPhysicalSlot* slot : ArrayRef(reusedSlots).take_front(kSummaryReuseLimit)) {
       uint64_t slotLogicalBytes = getSlotLogicalBytes(*slot, intervals);
-      os << "  slot #" << slot->id
+      os << "  slot #" << slot->id << "  addr=" << slot->address << "  size=" << formatReportMemory(slot->size)
-         << "  addr=" << slot->address
+         << "  intervals=" << slot->intervalIndices.size() << "  logical_sum=" << formatReportMemory(slotLogicalBytes)
-         << "  size=" << formatReportMemory(slot->size)
+         << "\n";
-         << "  intervals=" << slot->intervalIndices.size()
-         << "  logical_sum=" << formatReportMemory(slotLogicalBytes) << "\n";
       for (size_t intervalIndex : slot->intervalIndices) {
         const LocalAllocInterval& interval = intervals[intervalIndex];
-        os << "    #" << interval.id
+        os << "    #" << interval.id << "  [" << interval.start << "," << interval.end << "]"
-           << "  [" << interval.start << "," << interval.end << "]"
            << "  logical=" << formatReportMemory(interval.size)
            << "  first=" << summarizeOperation(interval.firstTouchOp, 40)
            << "  last=" << summarizeOperation(interval.lastTouchOp, 40) << "\n";
@@ -632,8 +629,7 @@ MemoryPlanArtifacts onnx_mlir::buildMemoryPlanArtifacts(Operation *coreLikeOp,
     const LocalAllocInterval& interval = intervals[slot->intervalIndices.front()];
     printedAttention = true;
     os << "  slot #" << slot->id << " is single-use"
-       << "  size=" << formatReportMemory(slot->size)
+       << "  size=" << formatReportMemory(slot->size) << "  interval=#" << interval.id
-       << "  interval=#" << interval.id
        << "  value=" << summarizeValue(interval.key.value, 56) << "\n";
     os << "    first=" << summarizeOperation(interval.firstTouchOp, 40)
        << "  last=" << summarizeOperation(interval.lastTouchOp, 40)
@@ -646,8 +642,7 @@ MemoryPlanArtifacts onnx_mlir::buildMemoryPlanArtifacts(Operation *coreLikeOp,
       continue;
     printedAttention = true;
     ++fallbackPrinted;
-    os << "  fallback interval #" << interval.id
+    os << "  fallback interval #" << interval.id << "  size=" << formatReportMemory(interval.size)
-       << "  size=" << formatReportMemory(interval.size)
        << "  value=" << summarizeValue(interval.key.value, 56) << "\n";
     os << "    reason: " << (interval.fallbackReason.empty() ? "<none>" : interval.fallbackReason) << "\n";
   }
@@ -659,10 +654,9 @@ MemoryPlanArtifacts onnx_mlir::buildMemoryPlanArtifacts(Operation *coreLikeOp,
       continue;
     printedAttention = true;
     ++nestedPrinted;
-    os << "  nested single-use interval #" << interval.id
+    os << "  nested single-use interval #" << interval.id << "  slot #" << interval.physicalSlotId
-       << "  slot #" << interval.physicalSlotId
+       << "  size=" << formatReportMemory(interval.size) << "  value=" << summarizeValue(interval.key.value, 56)
-       << "  size=" << formatReportMemory(interval.size)
+       << "\n";
-       << "  value=" << summarizeValue(interval.key.value, 56) << "\n";
     os << "    hint: move or sink this alloc inside the nested region if the IR allows it.\n";
   }
   if (!printedAttention)
@@ -674,14 +668,13 @@ MemoryPlanArtifacts onnx_mlir::buildMemoryPlanArtifacts(Operation *coreLikeOp,
       uint64_t slotLogicalBytes = getSlotLogicalBytes(slot, intervals);
       os << "  slot #" << slot.id << "  addr=" << slot.address << "  size=" << formatReportMemory(slot.size) << " ("
          << slot.size << ")"
-         << "  intervals=" << slot.intervalIndices.size()
+         << "  intervals=" << slot.intervalIndices.size() << "  logical_sum=" << formatReportMemory(slotLogicalBytes)
-         << "  logical_sum=" << formatReportMemory(slotLogicalBytes) << "\n";
+         << "\n";
       for (size_t intervalIndex : slot.intervalIndices) {
         const LocalAllocInterval& interval = intervals[intervalIndex];
         mlir::Value allocValue = interval.key.value;
         os << "    [" << interval.start << "," << interval.end << "]"
-           << "  #" << interval.id
+           << "  #" << interval.id << "  logical=" << formatReportMemory(interval.size)
-           << "  logical=" << formatReportMemory(interval.size)
            << "  nested=" << (interval.insideNestedRegion ? "yes" : "no")
            << "  escapes_loop=" << (interval.escapesLoop ? "yes" : "no")
            << "  first=" << summarizeOperation(interval.firstTouchOp, 48)
@@ -697,12 +690,10 @@ MemoryPlanArtifacts onnx_mlir::buildMemoryPlanArtifacts(Operation *coreLikeOp,
       const PlannedPhysicalSlot& slot = slots[interval.slotPlanIndex];
       mlir::Value allocValue = interval.key.value;
       Operation* definingOp = allocValue.getDefiningOp();
-      os << "  #" << interval.id
+      os << "  #" << interval.id << "  slot=" << slot.id << "  live=[" << interval.start << "," << interval.end << "]"
-         << "  slot=" << slot.id
-         << "  live=[" << interval.start << "," << interval.end << "]"
          << "  logical=" << formatReportMemory(interval.size)
-         << "  slot_size=" << formatReportMemory(interval.physicalSlotSize)
+         << "  slot_size=" << formatReportMemory(interval.physicalSlotSize) << "  addr=" << interval.assignedAddress
-         << "  addr=" << interval.assignedAddress << "\n";
+         << "\n";
       os << "    value=" << summarizeValue(allocValue, 88) << "\n";
       os << "    type=" << allocValue.getType() << "\n";
       os << "    loc="

src/PIM/Compiler/PimWeightEmitter.cpp

@@ -19,8 +19,7 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace {} // namespace
 
-WeightEmissionResult
+WeightEmissionResult createAndPopulateWeightFolder(ArrayRef<WeightFileRequest> requests, StringRef outputDirPath) {
-createAndPopulateWeightFolder(ArrayRef<WeightFileRequest> requests, StringRef outputDirPath) {
   auto coreWeightsDirPath = outputDirPath + "/weights";
   auto error = sys::fs::create_directory(coreWeightsDirPath);
   assert(!error && "Error creating weights directory");

src/PIM/Compiler/PimWeightEmitter.hpp

@@ -23,7 +23,7 @@ struct WeightEmissionResult {
   uint64_t totalWeightBytes = 0;
 };
 
-WeightEmissionResult
+WeightEmissionResult createAndPopulateWeightFolder(llvm::ArrayRef<WeightFileRequest> requests,
-createAndPopulateWeightFolder(llvm::ArrayRef<WeightFileRequest> requests, llvm::StringRef outputDirPath);
+                                                   llvm::StringRef outputDirPath);
 
 } // namespace onnx_mlir

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializeMergeSchedule.cpp

@@ -448,7 +448,10 @@ collectProducerKeysForDestinations(Value value, std::optional<ComputeInstance> l
     auto result = dyn_cast<OpResult>(value);
     if (!result)
       return {};
-    keys.push_back({{compute.getOperation(), 0, 1}, result.getResultNumber()});
+    keys.push_back({
+      {compute.getOperation(), 0, 1},
+      result.getResultNumber()
+    });
     return keys;
   }
 
@@ -476,8 +479,8 @@ collectProducerKeysForDestinations(Value value, std::optional<ComputeInstance> l
   return keys;
 }
 
-std::optional<ProducerKey>
+std::optional<ProducerKey> getInputRequestProducerKey(Value value,
-getInputRequestProducerKey(Value value, std::optional<ComputeInstance> logicalConsumer = std::nullopt) {
+                                                      std::optional<ComputeInstance> logicalConsumer = std::nullopt) {
   // Input resolution may request a whole-batch key for scalar consumers that read
   // a complete resultful compute_batch value.
   Operation* definingOp = value.getDefiningOp();
@@ -511,7 +514,10 @@ getInputRequestProducerKey(Value value, std::optional<ComputeInstance> logicalCo
     auto result = dyn_cast<OpResult>(value);
     if (!result)
       return std::nullopt;
-    return ProducerKey {{compute.getOperation(), 0, 1}, result.getResultNumber()};
+    return ProducerKey {
+      {compute.getOperation(), 0, 1},
+      result.getResultNumber()
+    };
   }
 
   if (auto batch = dyn_cast<SpatComputeBatch>(definingOp)) {
@@ -3268,9 +3274,7 @@ FailureOr<Value> materializeIndexedBatchRunReceive(MaterializerState& state,
 }
 
 FailureOr<SmallVector<Value, 4>>
-cloneInstanceBody(MaterializerState& state,
+cloneInstanceBody(MaterializerState& state, MaterializedClass& targetClass, ArrayRef<ComputeInstance> peers) {
-                  MaterializedClass& targetClass,
-                  ArrayRef<ComputeInstance> peers) {
   assert(!peers.empty() && "expected at least one peer instance");
   const ComputeInstance& instance = peers.front();
   Operation* sourceOp = instance.op;
@@ -3620,8 +3624,7 @@ FailureOr<SmallVector<Value, 4>> materializeBatchOutputGroupLoop(MaterializerSta
 
     state.rewriter.setInsertionPoint(targetClass.body->getTerminator());
     Value laneValue = getOrCreateIndexConstant(state.constantFolder, targetClass.op, item.laneStart);
-    return cloneBatchBodyForLane(
+    return cloneBatchBodyForLane(state, targetClass, item, laneValue, group.resultIndices, {});
-      state, targetClass, item, laneValue, group.resultIndices, {});
   }
 
   state.rewriter.setInsertionPoint(targetClass.body->getTerminator());
@@ -3733,8 +3736,7 @@ FailureOr<MaterializationRun> collectBatchMaterializationRun(MaterializerState&
     if (state.materializedLogicalSlots.contains(classSlot))
       break;
 
-    FailureOr<SmallVector<ComputeInstance, 8>> peers =
+    FailureOr<SmallVector<ComputeInstance, 8>> peers = getMaterializationRunSlotPeers(state, targetClass, slot);
-      getMaterializationRunSlotPeers(state, targetClass, slot);
     if (failed(peers) || peers->empty())
       break;
 
@@ -3818,12 +3820,11 @@ bool hasMaterializationRunGroupSameClassConsumer(MaterializerState& state,
                                                  const OutputDestinationGroup& group) {
   for (size_t resultIndex : group.resultIndices) {
     for (const MaterializationRunSlot& slot : run) {
-      for (const ComputeInstance& peer : slot.peers) {
+      for (const ComputeInstance& peer : slot.peers)
         if (hasSameClassConsumer(state, {peer, resultIndex}, classId))
           return true;
     }
   }
-  }
 
   return false;
 }
@@ -4236,7 +4237,8 @@ LogicalResult materializeInstanceSlot(MaterializerState& state, const ComputeIns
     return success();
 
   if (isa<SpatComputeBatch>(instance.op)) {
-    FailureOr<MaterializationRun> run = collectBatchMaterializationRun(state, targetClass, startLogicalSlot, instance.op);
+    FailureOr<MaterializationRun> run =
+      collectBatchMaterializationRun(state, targetClass, startLogicalSlot, instance.op);
 
     if (succeeded(run)) {
       if (!targetClass.isBatch)

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp

@@ -197,16 +197,13 @@ SmallVector<ProducerValueRef, 4> collectProducerValueRefs(Value value, const Com
         return producers;
       }
 
-      if (isa<SpatComputeBatch>(consumerInstance.op)) {
+      if (isa<SpatComputeBatch>(consumerInstance.op))
         for (ComputeInstance instance :
              getBatchChunksForRange(batch, consumerInstance.laneStart, consumerInstance.laneCount))
           producers.push_back({instance, 0});
-      }
+      else
-      else {
+        for (ComputeInstance instance : getBatchChunksForRange(batch, 0, static_cast<uint32_t>(batch.getLaneCount())))
-        for (ComputeInstance instance :
-             getBatchChunksForRange(batch, 0, static_cast<uint32_t>(batch.getLaneCount())))
           producers.push_back({instance, 0});
-      }
       return producers;
     }
 
@@ -217,16 +214,17 @@ SmallVector<ProducerValueRef, 4> collectProducerValueRefs(Value value, const Com
   }
 
   if (auto compute = dyn_cast<SpatCompute>(op)) {
-    producers.push_back({ComputeInstance {compute.getOperation(), 0, 1},
+    producers.push_back({
-                         static_cast<size_t>(cast<OpResult>(value).getResultNumber())});
+      ComputeInstance {compute.getOperation(), 0, 1},
+       static_cast<size_t>(cast<OpResult>(value).getResultNumber())
+    });
     return producers;
   }
 
   if (auto batch = dyn_cast<SpatComputeBatch>(op)) {
     if (batch.getNumResults() != 0) {
       uint32_t laneStart = isa<SpatComputeBatch>(consumerInstance.op) ? consumerInstance.laneStart : 0;
-      uint32_t laneCount = isa<SpatComputeBatch>(consumerInstance.op)
+      uint32_t laneCount = isa<SpatComputeBatch>(consumerInstance.op) ? consumerInstance.laneCount
-                             ? consumerInstance.laneCount
                                                                       : static_cast<uint32_t>(batch.getLaneCount());
       for (ComputeInstance instance : getBatchChunksForRange(batch, laneStart, laneCount))
         producers.push_back({instance, 0});
@@ -242,7 +240,9 @@ SmallVector<ProducerValueRef, 4> collectProducerValueRefs(Value value, const Com
   return producers;
 }
 
-Cost getProducerTransferCost(Value input, const ComputeInstance& consumerInstance, const ProducerValueRef& producerRef) {
+Cost getProducerTransferCost(Value input,
+                             const ComputeInstance& consumerInstance,
+                             const ProducerValueRef& producerRef) {
   Cost transferCost = getInputTransferCost(consumerInstance, input);
   auto producerBatch = dyn_cast<SpatComputeBatch>(producerRef.instance.op);
   if (!producerBatch || producerBatch.getNumResults() == 0)
@@ -256,9 +256,8 @@ Cost getProducerTransferCost(Value input, const ComputeInstance& consumerInstanc
     }
   }
 
-  return scaleTransferCostByLaneCount(transferCost,
+  return scaleTransferCostByLaneCount(
-                                      static_cast<uint32_t>(producerBatch.getLaneCount()),
+    transferCost, static_cast<uint32_t>(producerBatch.getLaneCount()), producerRef.instance.laneCount);
-                                      producerRef.instance.laneCount);
 }
 
 static CrossbarWeight getOpaqueCrossbarWeight(Value value, std::optional<uint32_t> lane) {

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp

@@ -64,9 +64,8 @@ ComputeInstance getBatchChunkForLane(SpatComputeBatch batch, uint32_t lane) {
   return getBatchChunkForIndex(batch, getBatchChunkIndexForLane(batch.getLaneCount(), lane));
 }
 
-llvm::SmallVector<ComputeInstance, 4> getBatchChunksForRange(SpatComputeBatch batch,
+llvm::SmallVector<ComputeInstance, 4>
-                                                             uint32_t laneStart,
+getBatchChunksForRange(SpatComputeBatch batch, uint32_t laneStart, uint32_t laneCount) {
-                                                             uint32_t laneCount) {
   llvm::SmallVector<ComputeInstance, 4> chunks;
   if (laneCount == 0)
     return chunks;

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.hpp

@@ -32,9 +32,8 @@ BatchChunkRange getBatchChunkRange(int32_t laneCount, size_t chunkIndex);
 size_t getBatchChunkIndexForLane(int32_t laneCount, uint32_t lane);
 ComputeInstance getBatchChunkForIndex(SpatComputeBatch batch, size_t chunkIndex);
 ComputeInstance getBatchChunkForLane(SpatComputeBatch batch, uint32_t lane);
-llvm::SmallVector<ComputeInstance, 4> getBatchChunksForRange(SpatComputeBatch batch,
+llvm::SmallVector<ComputeInstance, 4>
-                                                             uint32_t laneStart,
+getBatchChunksForRange(SpatComputeBatch batch, uint32_t laneStart, uint32_t laneCount);
-                                                             uint32_t laneCount);
 
 std::optional<ProducerValueRef> getProducerValueRef(mlir::Value value,
                                                     const ComputeInstance* consumerInstance = nullptr);