A new Beginning

src/PIM/Common/Support/DebugDump.cpp

@@ -17,7 +17,7 @@ std::fstream openDialectDumpFileWithExtension(const std::string& name, llvm::Str
   return std::fstream(dialectsDir + "/" + name + "." + extension.str(), std::ios::out);
 }
 
-void dumpModule(mlir::ModuleOp moduleOp, const std::string& name) {
+void dumpModule(mlir::ModuleOp moduleOp, const std::string& name, bool assumeVerified) {
   std::fstream file = openDialectDumpFileWithExtension(name, "/dialects", "mlir");
   if (!file.is_open())
     return;
@@ -25,6 +25,8 @@ void dumpModule(mlir::ModuleOp moduleOp, const std::string& name) {
   llvm::raw_os_ostream os(file);
   mlir::OpPrintingFlags flags;
   flags.elideLargeElementsAttrs().enableDebugInfo(true, false);
+  if (assumeVerified)
+    flags.assumeVerified();
   moduleOp.print(os, flags);
   os.flush();
   file.close();

src/PIM/Common/Support/DebugDump.hpp

@@ -10,7 +10,7 @@ namespace onnx_mlir {
 
 /// Emits a MLIR snapshot under the current compiler output
 /// directory for pass-level debugging.
-void dumpModule(mlir::ModuleOp moduleOp, const std::string& name);
+void dumpModule(mlir::ModuleOp moduleOp, const std::string& name, bool assumeVerified = false);
 
 /// Opens a file under the same dialect dump directory used by dumpModule.
 std::fstream openDialectDumpFileWithExtension(const std::string& name,llvm::StringRef destination = "/dialects",  llvm::StringRef extension = "mlir");

src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp

@@ -18,7 +18,6 @@
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Accelerators/PIM/Pass/PIMPasses.h"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -380,12 +379,6 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
       signalPassFailure();
     } else {
       dumpModule(moduleOp, "spatial1_graph");
-      spatial::SpatialDataflowExportStage exportMode = spatial::getSpatialDataflowExportStage();
-      if (spatial::shouldExportSpatialDataflowStage(exportMode, spatial::SpatialDataflowExportStage::Pre)
-          && failed(spatial::exportSpatialDataflowCsvPre(funcOp))) {
-        signalPassFailure();
-        return;
-      }
     }
 
     if (!verifyLogicalPhase("at the end of LowerSpatialPlans"))

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp

@@ -56,13 +56,18 @@ bool isLegalExternalCapture(Value value, Region& region) {
   return definingOp && definingOp->hasTrait<OpTrait::ConstantLike>();
 }
 
+bool isRecordedDeferredCommunicationSource(Operation* op, Value value) {
+  auto transfer = dyn_cast<spatial::SpatDeferredCommunicationOp>(op);
+  return transfer && llvm::is_contained(transfer.getSources(), value);
+}
+
 template <typename ComputeOpTy>
 void verifyComputeBodyCaptures(ComputeOpTy compute, StringRef kind, pim::CappedDiagnosticReporter& diagnostics) {
   Region& body = compute.getBody();
   body.walk([&](Operation* nestedOp) {
     for (OpOperand& operand : nestedOp->getOpOperands()) {
       Value value = operand.get();
-      if (isLegalExternalCapture(value, body))
+      if (isLegalExternalCapture(value, body) || isRecordedDeferredCommunicationSource(nestedOp, value))
         continue;
 
       Operation* definingOp = value.getDefiningOp();
@@ -90,21 +95,29 @@ bool isLegalHostBackedValue(Value value) {
   return definingOp->getDialect()->getNamespace() != "spat";
 }
 
+bool isScheduledPhase1Value(Value value) {
+  Operation* definingOp = value.getDefiningOp();
+  return isa_and_nonnull<spatial::SpatScheduledCompute, spatial::SpatScheduledComputeBatch>(definingOp);
+}
+
 template <typename ComputeOpTy>
 void verifyScheduledInputs(ComputeOpTy compute,
                            bool allowChannelReceiveInputs,
                            StringRef kind,
                            pim::CappedDiagnosticReporter& diagnostics) {
   for (auto [inputIndex, input] : llvm::enumerate(compute.getInputs())) {
+    size_t currentInputIndex = inputIndex;
     Operation* definingOp = input.getDefiningOp();
     if (allowChannelReceiveInputs && isa_and_nonnull<spatial::SpatChannelReceiveOp>(definingOp))
       continue;
+    if (isScheduledPhase1Value(input))
+      continue;
     if (isLegalHostBackedValue(input))
       continue;
 
     diagnostics.report(compute.getOperation(), [&](Operation* illegalOp) {
       InFlightDiagnostic diag = illegalOp->emitOpError()
-                                << kPhaseMarker << " " << kind << " input #" << inputIndex
+                                << kPhaseMarker << " " << kind << " input #" << currentInputIndex
                                 << (allowChannelReceiveInputs ? " must come from the host or explicit spat.channel_receive"
                                                               : " must come from the host");
       if (definingOp)
@@ -163,9 +176,9 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter
 
 void verifyScheduledTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter& diagnostics) {
   for (Operation& op : funcOp.getOps()) {
-    if (isa<spatial::SpatGraphCompute, spatial::SpatGraphComputeBatch>(&op)) {
+    if (isa<spatial::SpatChannelSendOp, spatial::SpatChannelReceiveOp>(&op)) {
       diagnostics.report(&op, [&](Operation* illegalOp) {
-        illegalOp->emitOpError() << kPhaseMarker << " graph Spatial compute op remained after merge materialization";
+        illegalOp->emitOpError() << kPhaseMarker << " real channel communication is not allowed in scheduled phase 1";
       });
     }
   }

src/PIM/Dialect/Spatial/CMakeLists.txt

@@ -7,14 +7,9 @@ add_pim_library(SpatialOps
   SpatialOpsVerify.cpp
   SpatialOpsCanonicalization.cpp
   ${PIM_SRC_ROOT}/Conversion/ONNXToSpatial/CompileTime.cpp
-  Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp
-  Transforms/MergeComputeNodes/CommunicationPlan.cpp
-  Transforms/MergeComputeNodes/HostOutputFinalization.cpp
-  Transforms/MergeComputeNodes/MaterializeMergeSchedule.cpp
-  Transforms/MergeComputeNodes/ProjectedFragments.cpp
-  Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp
   Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp
   Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp
+  Transforms/MergeComputeNodes/Scheduling/MergeComputeNodesPass.cpp
   Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp
   Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp
 

src/PIM/Dialect/Spatial/Spatial.td

@@ -6,6 +6,7 @@ include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/BuiltinTypes.td"
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/RegionKindInterface.td"
+include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/ParallelCombiningOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
@@ -27,7 +28,7 @@ def SpatTensor :
 //===----------------------------------------------------------------------===//
 
 class SpatComputeLikeBase<string mnemonic> : SpatOp<mnemonic,
-    [SingleBlock, AttrSizedOperandSegments,
+    [AttrSizedOperandSegments,
      DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmBlockArgumentNames"]>]> {
   let summary = "Compute region with attached constant weights";
 
@@ -40,7 +41,7 @@ class SpatComputeLikeBase<string mnemonic> : SpatOp<mnemonic,
     Variadic<SpatTensor>:$outputs
   );
 
-  let regions = (region SizedRegion<1>:$body);
+  let regions = (region MinSizedRegion<1>:$body);
 
   let hasVerifier = 1;
   let hasFolder = 1;
@@ -76,7 +77,7 @@ def SpatScheduledCompute : SpatComputeLikeBase<"scheduled_compute"> {
 }
 
 class SpatComputeBatchLikeBase<string mnemonic> : SpatOp<mnemonic,
-    [SingleBlock, AttrSizedOperandSegments,
+    [AttrSizedOperandSegments,
      DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmBlockArgumentNames"]>]> {
   let summary = "Tensor-native batch of equivalent compute lanes with shared weights and packed inputs";
 
@@ -90,7 +91,7 @@ class SpatComputeBatchLikeBase<string mnemonic> : SpatOp<mnemonic,
     Variadic<SpatTensor>:$outputs
   );
 
-  let regions = (region SizedRegion<1>:$body);
+  let regions = (region MinSizedRegion<1>:$body);
 
   let hasVerifier = 1;
   let hasCustomAssemblyFormat = 1;
@@ -161,6 +162,54 @@ def SpatYieldOp : SpatOp<"yield", [Terminator]> {
   let hasCustomAssemblyFormat = 1;
 }
 
+def SpatBlockYieldOp : SpatOp<"block_yield", [
+    Terminator,
+    DeclareOpInterfaceMethods<BranchOpInterface, ["getSuccessorForOperands"]>
+  ]> {
+  let summary = "Terminate a scheduled structural compute block";
+
+  let arguments = (ins
+    Variadic<AnyType>:$outputs
+  );
+
+  let successors = (successor
+    VariadicSuccessor<AnySuccessor>:$next
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def SpatDeferredCommunicationOp : SpatOp<"deferred_communication", [SingleBlock]> {
+  let summary = "Temporary scheduled communication placeholder";
+
+  let arguments = (ins
+    Variadic<SpatTensor>:$sources,
+    StrAttr:$exchangeId,
+    StrAttr:$logicalProducer,
+    StrAttr:$logicalConsumer,
+    I64Attr:$sourceClass,
+    I64Attr:$targetClass,
+    I64Attr:$sourceCore,
+    I64Attr:$targetCore,
+    I64Attr:$sourceLane,
+    I64Attr:$targetLane,
+    StrAttr:$transferKind,
+    I64Attr:$resultIndex,
+    DenseI64ArrayAttr:$projectedTransfer,
+    I64Attr:$hostOutputOwner
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let regions = (region SizedRegion<1>:$body);
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
 def SpatExtractRowsOp : SpatOp<"extract_rows", []> {
   let summary = "Extract every row of a rank-2 tensor as separate rank-2 row tensors";
 

src/PIM/Dialect/Spatial/SpatialOps.cpp

@@ -238,6 +238,15 @@ void SpatScheduledCompute::getAsmBlockArgumentNames(Region& region, OpAsmSetValu
   setComputeAsmBlockArgumentNames(*this, region, setNameFn);
 }
 
+SuccessorOperands SpatBlockYieldOp::getSuccessorOperands(unsigned index) {
+  assert(index == 0 && "invalid successor index");
+  return SuccessorOperands(getOutputsMutable());
+}
+
+Block* SpatBlockYieldOp::getSuccessorForOperands(ArrayRef<Attribute>) {
+  return getOperation()->getNumSuccessors() == 0 ? nullptr : getOperation()->getSuccessor(0);
+}
+
 std::optional<BlockArgument> SpatGraphComputeBatch::getLaneArgument() { return getBlockArgument(getBody(), 0); }
 std::optional<BlockArgument> SpatGraphComputeBatch::getWeightArgument(unsigned idx) {
   return getBlockArgument(getBody(), 1 + idx);

src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp

@@ -50,6 +50,64 @@ static void printBlockArgumentList(OpAsmPrinter& printer, ArrayRef<BlockArgument
   printer << ")";
 }
 
+static void printBlockHeaderWithoutArgLocs(OpAsmPrinter& printer, Block& block) {
+  printer.printSuccessor(&block);
+  if (block.getNumArguments() == 0) {
+    printer << ":";
+    return;
+  }
+
+  printer << "(";
+  for (auto [index, argument] : llvm::enumerate(block.getArguments())) {
+    if (index != 0)
+      printer << ", ";
+    printer.printOperand(argument);
+    printer << ": ";
+    printer.printType(argument.getType());
+  }
+  printer << "):";
+}
+
+static void printRegionWithoutBlockArgLocs(OpAsmPrinter& printer,
+                                           Region& region,
+                                           bool printEntryBlockArgs = true,
+                                           bool printBlockTerminators = true,
+                                           bool printEmptyBlock = false,
+                                           bool printEntryBlockHeaderWhenMultiblock = false) {
+  printer << " {";
+  if (region.empty()) {
+    if (printEmptyBlock)
+      printer << "\n";
+    printer << "}";
+    return;
+  }
+
+  bool hasMultipleBlocks = std::next(region.begin()) != region.end();
+  for (auto [blockIndex, block] : llvm::enumerate(region)) {
+    bool printBlockHeader = blockIndex != 0 || printEntryBlockArgs
+                            || (blockIndex == 0 && printEntryBlockHeaderWhenMultiblock && hasMultipleBlocks)
+                            || (printEmptyBlock && block.empty());
+    unsigned indent = printBlockHeader ? 4u : 2u;
+
+    if (printBlockHeader) {
+      printer.getStream() << "\n";
+      printer.getStream().indent(2);
+      printBlockHeaderWithoutArgLocs(printer, block);
+    }
+
+    for (Operation& nestedOp : block) {
+      if (!printBlockTerminators && nestedOp.hasTrait<OpTrait::IsTerminator>())
+        continue;
+      printer.getStream() << "\n";
+      printer.getStream().indent(indent);
+      printer.printCustomOrGenericOp(&nestedOp);
+    }
+  }
+
+  printer.getStream() << "\n";
+  printer << "}";
+}
+
 static ParseResult parseBlockArgumentList(OpAsmParser& parser, SmallVectorImpl<OpAsmParser::Argument>& arguments) {
   if (parser.parseLParen())
     return failure();
@@ -160,7 +218,9 @@ void printComputeLikeOp(ComputeOpTy op, OpAsmPrinter& printer) {
   printer << " -> ";
   printCompressedTypeSequence(printer, op.getResultTypes());
   printer << " ";
-  printer.printRegion(op.getBody(), /*printEntryBlockArgs=*/false);
+  printRegionWithoutBlockArgLocs(
+    printer, op.getBody(), /*printEntryBlockArgs=*/false, /*printBlockTerminators=*/true,
+    /*printEmptyBlock=*/false, /*printEntryBlockHeaderWhenMultiblock=*/true);
 }
 
 template <typename ComputeOpTy>
@@ -290,7 +350,9 @@ void printComputeBatchLikeOp(ComputeBatchOpTy op, OpAsmPrinter& printer) {
   printer << " -> ";
   printCompressedTypeSequence(printer, op.getResultTypes());
   printer << " ";
-  printer.printRegion(op.getBody(), /*printEntryBlockArgs=*/false);
+  printRegionWithoutBlockArgLocs(
+    printer, op.getBody(), /*printEntryBlockArgs=*/false, /*printBlockTerminators=*/true,
+    /*printEmptyBlock=*/false, /*printEntryBlockHeaderWhenMultiblock=*/true);
 }
 
 template <typename ComputeBatchOpTy>
@@ -407,6 +469,89 @@ ParseResult SpatYieldOp::parse(OpAsmParser& parser, OperationState& result) {
   return parser.resolveOperands(outputs, outputTypes, parser.getCurrentLocation(), result.operands);
 }
 
+void SpatBlockYieldOp::print(OpAsmPrinter& printer) {
+  printer << " ";
+  printCompressedValueSequence(printer, getOutputs());
+  if (getOperation()->getNumSuccessors() != 0) {
+    printer << " next ";
+    printer.printSuccessor(getOperation()->getSuccessor(0));
+  }
+  printer.printOptionalAttrDict((*this)->getAttrs());
+  printer << " : ";
+  printCompressedTypeSequence(printer, getOutputs().getTypes());
+}
+
+ParseResult SpatBlockYieldOp::parse(OpAsmParser& parser, OperationState& result) {
+  SmallVector<OpAsmParser::UnresolvedOperand> outputs;
+  SmallVector<Type> outputTypes;
+  Block* successor = nullptr;
+
+  OpAsmParser::UnresolvedOperand firstOutput;
+  OptionalParseResult firstOutputResult = parser.parseOptionalOperand(firstOutput);
+  if (firstOutputResult.has_value()) {
+    if (failed(*firstOutputResult))
+      return failure();
+    if (parseCompressedOperandEntryWithFirst(parser, firstOutput, outputs))
+      return failure();
+    while (succeeded(parser.parseOptionalComma()))
+      if (parseOneCompressedOperandEntry(parser, outputs))
+        return failure();
+  }
+
+  if (succeeded(parser.parseOptionalKeyword("next")) && parser.parseSuccessor(successor))
+    return failure();
+
+  if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()
+      || parseCompressedTypeSequence(parser, outputTypes, /*allowEmpty=*/true))
+    return failure();
+
+  if (outputs.size() != outputTypes.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of outputs and output types must match");
+  if (parser.resolveOperands(outputs, outputTypes, parser.getCurrentLocation(), result.operands))
+    return failure();
+  if (successor)
+    result.addSuccessors(successor);
+  return success();
+}
+
+void SpatDeferredCommunicationOp::print(OpAsmPrinter& printer) {
+  printer << " ";
+  printCompressedValueSequence(printer, getSources());
+  printer.printOptionalAttrDict((*this)->getAttrs());
+  printer << " : ";
+  printer.printFunctionalType(getSources().getTypes(), getOperation()->getResultTypes());
+  printer << " ";
+  printRegionWithoutBlockArgLocs(printer, getBody(), /*printEntryBlockArgs=*/false);
+}
+
+ParseResult SpatDeferredCommunicationOp::parse(OpAsmParser& parser, OperationState& result) {
+  SmallVector<OpAsmParser::UnresolvedOperand> sources;
+  Type functionTypeStorage;
+
+  if (parseCompressedOperandSequence(parser, sources) || parser.parseOptionalAttrDict(result.attributes)
+      || parser.parseColon() || parser.parseType(functionTypeStorage))
+    return failure();
+
+  auto functionType = dyn_cast<FunctionType>(functionTypeStorage);
+  if (!functionType)
+    return parser.emitError(parser.getCurrentLocation(), "expected deferred communication function type");
+  if (sources.size() != functionType.getNumInputs())
+    return parser.emitError(parser.getCurrentLocation(), "number of sources and source types must match");
+
+  if (parser.resolveOperands(sources, functionType.getInputs(), parser.getCurrentLocation(), result.operands))
+    return failure();
+  result.addTypes(functionType.getResults());
+
+  Region* body = result.addRegion();
+  SmallVector<OpAsmParser::Argument> bodyArgs;
+  for (Type type : functionType.getInputs()) {
+    OpAsmParser::Argument argument;
+    argument.type = type;
+    bodyArgs.push_back(argument);
+  }
+  return parser.parseRegion(*body, bodyArgs);
+}
+
 void SpatExtractRowsOp::print(OpAsmPrinter& printer) {
   printer << " ";
   printer.printOperand(getInput());
@@ -618,7 +763,8 @@ ParseResult SpatScheduledComputeBatch::parse(OpAsmParser& parser, OperationState
 
 void SpatInParallelOp::print(OpAsmPrinter& printer) {
   printer << " ";
-  printer.printRegion(getRegion(), /*printEntryBlockArgs=*/false, /*printBlockTerminators=*/false);
+  printRegionWithoutBlockArgLocs(
+    printer, getRegion(), /*printEntryBlockArgs=*/false, /*printBlockTerminators=*/false);
   printer.printOptionalAttrDict((*this)->getAttrs());
 }
 

src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp

@@ -192,12 +192,18 @@ static bool isConstantExternalValue(Value value) {
   return definingOp && definingOp->hasTrait<OpTrait::ConstantLike>();
 }
 
+static bool isRecordedDeferredCommunicationSource(Operation* op, Value value) {
+  auto transfer = dyn_cast<SpatDeferredCommunicationOp>(op);
+  return transfer && llvm::is_contained(transfer.getSources(), value);
+}
+
 static LogicalResult verifyOnlyConstantExternalValues(Operation* ownerOp, Region& region, StringRef kind) {
   bool hasFailure = false;
   region.walk([&](Operation* op) {
     for (OpOperand& operand : op->getOpOperands()) {
       Value value = operand.get();
-      if (isDefinedInsideRegion(value, region) || isConstantExternalValue(value))
+      if (isDefinedInsideRegion(value, region) || isConstantExternalValue(value)
+          || isRecordedDeferredCommunicationSource(op, value))
         continue;
 
       InFlightDiagnostic diagnostic =
@@ -219,8 +225,35 @@ static LogicalResult verifyOnlyConstantExternalValues(Operation* ownerOp, Region
   return success(!hasFailure);
 }
 
+static LogicalResult verifyYieldTypes(Operation* op, Region& region, TypeRange resultTypes, StringRef kind) {
+  if (region.empty())
+    return op->emitOpError() << kind << " requires a body block";
+  Block& block = region.front();
+  auto yield = dyn_cast_or_null<SpatYieldOp>(block.getTerminator());
+  if (!yield)
+    return op->emitOpError() << kind << " body must terminate with spat.yield";
+  if (yield.getOutputs().size() != resultTypes.size())
+    return op->emitOpError() << kind << " yield operand count must match result count";
+  for (auto [yieldType, resultType] : llvm::zip(yield.getOutputs().getTypes(), resultTypes))
+    if (yieldType != resultType)
+      return op->emitOpError() << kind << " yield operand types must match result types";
+  return success();
+}
+
+static LogicalResult verifyRegionArguments(Operation* op, Region& region, ValueRange operands, StringRef kind) {
+  if (region.empty())
+    return op->emitOpError() << kind << " requires a body block";
+  Block& block = region.front();
+  if (block.getNumArguments() != operands.size())
+    return op->emitOpError() << kind << " body argument count must match operand count";
+  for (auto [arg, operand] : llvm::zip(block.getArguments(), operands))
+    if (arg.getType() != operand.getType())
+      return op->emitOpError() << kind << " body argument types must match operand types";
+  return success();
+}
+
 template <typename ComputeBatchOpTy>
-static LogicalResult verifyBatchBody(ComputeBatchOpTy batchOp, Block& block) {
+static LogicalResult verifyBatchBody(ComputeBatchOpTy batchOp, Block& block, bool verifyLaneSliceOffsets = true) {
   if (batchOp.getNumResults() == 0) {
     auto yieldOp = dyn_cast_or_null<SpatYieldOp>(block.getTerminator());
     if (!yieldOp)
@@ -235,11 +268,12 @@ static LogicalResult verifyBatchBody(ComputeBatchOpTy batchOp, Block& block) {
   auto laneArg = batchOp.getLaneArgument();
   if (!laneArg)
     return batchOp.emitError("compute_batch body must have a lane block argument");
-  for (auto& bodyOp : block) {
+  if (verifyLaneSliceOffsets)
-    if (auto extractSlice = dyn_cast<tensor::ExtractSliceOp>(&bodyOp))
+    for (auto& bodyOp : block) {
-      if (failed(verifyStaticUnitStrideExtractSliceOp(extractSlice, *laneArg, "tensor.extract_slice")))
+      if (auto extractSlice = dyn_cast<tensor::ExtractSliceOp>(&bodyOp))
-        return failure();
+        if (failed(verifyStaticUnitStrideExtractSliceOp(extractSlice, *laneArg, "tensor.extract_slice")))
-  }
+          return failure();
+    }
   return success();
 }
 
@@ -679,7 +713,9 @@ LogicalResult verifyComputeResultsUses(Operation* op) {
   if (!isAnySpatialComputeLike(op))
     return op->emitError("verifyComputeResultUses: op is not a Spatial compute-like operation");
   if (!llvm::all_of(op->getResults(), [](Value result) {
-        return llvm::all_of(result.getUsers(), [](Operation* op) {
+        return llvm::all_of(result.getUsers(), [result](Operation* op) {
+          if (isRecordedDeferredCommunicationSource(op, result))
+            return true;
           return !isAnySpatialComputeLike(op->getParentOp());
         });
       })) {
@@ -690,57 +726,67 @@ LogicalResult verifyComputeResultsUses(Operation* op) {
 
 template <typename ComputeOpTy>
 LogicalResult verifyComputeLikeOp(ComputeOpTy compute, StringRef opName) {
-  auto& block = compute.getBody().front();
   unsigned expectedArgCount = compute.getWeights().size() + compute.getInputs().size();
-  if (block.getNumArguments() != expectedArgCount)
+  bool isScheduled = isa<SpatScheduledCompute>(compute.getOperation());
-    return compute.emitOpError("compute body must have weight and input block arguments");
+  if (compute.getBody().empty())
+    return compute.emitOpError("compute body must have at least one block");
 
-  for (auto [weightIndex, weight] : llvm::enumerate(compute.getWeights())) {
+  SmallVector<Type> yieldedTypes;
-    auto blockArg = compute.getWeightArgument(weightIndex);
+  for (Block& block : compute.getBody()) {
-    if (!blockArg || blockArg->getType() != weight.getType())
+    if ((!isScheduled && block.getNumArguments() != expectedArgCount)
-      return compute.emitOpError("compute weight block argument types must match weight operand types exactly");
+        || (isScheduled && block.getNumArguments() < expectedArgCount))
-  }
+      return compute.emitOpError("compute body must have weight and input block arguments");
-  for (auto [inputIndex, input] : llvm::enumerate(compute.getInputs())) {
-    auto blockArg = compute.getInputArgument(inputIndex);
-    if (!blockArg || blockArg->getType() != input.getType())
-      return compute.emitOpError("compute input block argument types must match input operand types exactly");
-  }
 
-  if (block.mightHaveTerminator()) {
+    for (auto [weightIndex, weight] : llvm::enumerate(compute.getWeights()))
-    auto yieldOp = dyn_cast_or_null<SpatYieldOp>(block.getTerminator());
+      if (block.getArgument(weightIndex).getType() != weight.getType())
-    if (!yieldOp)
+        return compute.emitOpError("compute weight block argument types must match weight operand types exactly");
+    for (auto [inputIndex, input] : llvm::enumerate(compute.getInputs()))
+      if (block.getArgument(compute.getWeights().size() + inputIndex).getType() != input.getType())
+        return compute.emitOpError("compute input block argument types must match input operand types exactly");
+
+    Operation* terminator = block.getTerminator();
+    if (auto yieldOp = dyn_cast_or_null<SpatYieldOp>(terminator)) {
+      if (isScheduled)
+        return compute.emitOpError("scheduled compute blocks must terminate with spat.block_yield");
+      llvm::append_range(yieldedTypes, yieldOp->getOperandTypes());
+      continue;
+    }
+    auto blockYield = dyn_cast_or_null<SpatBlockYieldOp>(terminator);
+    if (!blockYield || !isScheduled)
       return compute.emitOpError("ComputeOp must have a single yield operation");
+    if (blockYield->getNumSuccessors() == 0)
+      llvm::append_range(yieldedTypes, blockYield->getOperandTypes());
+  }
 
-    auto resultTypes = compute.getResultTypes();
+  auto resultTypes = compute.getResultTypes();
-    auto yieldTypes = yieldOp->getOperandTypes();
+  if (resultTypes.size() != yieldedTypes.size())
-    if (resultTypes.size() != yieldTypes.size())
+    return compute.emitOpError("ComputeOp must have same number of results as yielded operands");
-      return compute.emitOpError("ComputeOp must have same number of results as yieldOp operands");
 
-    for (auto it : llvm::reverse(llvm::zip(resultTypes, yieldTypes))) {
+  for (auto it : llvm::reverse(llvm::zip(resultTypes, yieldedTypes))) {
-      auto resultType = std::get<0>(it);
+    auto resultType = std::get<0>(it);
-      auto yieldType = std::get<1>(it);
+    auto yieldType = std::get<1>(it);
 
-      if (resultType != yieldType || failed(verifyCompatibleShape(resultType, yieldType)))
+    if (resultType != yieldType || failed(verifyCompatibleShape(resultType, yieldType)))
-        return compute.emitOpError("ComputeOp output must be of the same type as yieldOp operand");
+      return compute.emitOpError("ComputeOp output must be of the same type as yieldOp operand");
 
-      if (auto resultRankedType = dyn_cast<RankedTensorType>(resultType)) {
+    if (auto resultRankedType = dyn_cast<RankedTensorType>(resultType)) {
-        if (auto yieldRankedType = dyn_cast<RankedTensorType>(yieldType)) {
+      if (auto yieldRankedType = dyn_cast<RankedTensorType>(yieldType)) {
-          if (resultRankedType.getEncoding() != yieldRankedType.getEncoding())
+        if (resultRankedType.getEncoding() != yieldRankedType.getEncoding())
-            return compute.emitOpError("ComputeOp output must have the same encoding as yieldOp operand");
-        }
-        else {
           return compute.emitOpError("ComputeOp output has an encoding while yieldOp operand does not have one");
-        }
       }
-      else if (dyn_cast<RankedTensorType>(yieldType)) {
+      else {
-        return compute.emitOpError("ComputeOp output must not have an encoding if yieldOp operand has one");
+        return compute.emitOpError("ComputeOp output must have the same encoding as yieldOp operand");
       }
     }
+    else if (dyn_cast<RankedTensorType>(yieldType)) {
+      return compute.emitOpError("ComputeOp output must not have an encoding if yieldOp operand has one");
+    }
   }
 
-  for (unsigned inputIndex = 0; inputIndex < compute.getInputs().size(); ++inputIndex)
+  if (compute.getBody().hasOneBlock())
-    if (auto inputArg = compute.getInputArgument(inputIndex); !inputArg || inputArg->use_empty())
+    for (unsigned inputIndex = 0; inputIndex < compute.getInputs().size(); ++inputIndex)
-      return compute.emitOpError("ComputeOp block argument is not used");
+      if (auto inputArg = compute.getInputArgument(inputIndex); !inputArg || inputArg->use_empty())
+        return compute.emitOpError("ComputeOp block argument is not used");
   if (failed(verifyStaticWeights(compute, opName)))
     return failure();
   if (failed(verifyOnlyConstantExternalValues(compute.getOperation(), compute.getBody(), opName)))
@@ -754,6 +800,31 @@ LogicalResult SpatGraphCompute::verify() { return verifyComputeLikeOp(*this, "sp
 
 LogicalResult SpatScheduledCompute::verify() { return verifyComputeLikeOp(*this, "spat.scheduled_compute"); }
 
+LogicalResult SpatBlockYieldOp::verify() {
+  if (getOperation()->getNumSuccessors() > 1)
+    return emitOpError("may target at most one next scheduled block");
+  Operation* parent = getOperation()->getParentOp();
+  if (!isa_and_nonnull<SpatScheduledCompute>(parent))
+    return emitOpError("expected spat.scheduled_compute parent");
+  if (getOperation()->getNumSuccessors() == 1) {
+    Block* next = getOperation()->getSuccessor(0);
+    if (getOperation()->getNumOperands() != next->getNumArguments())
+      return emitOpError("successor operand count must match next block argument count");
+    for (auto [operand, argument] : llvm::zip(getOperation()->getOperands(), next->getArguments()))
+      if (operand.getType() != argument.getType())
+        return emitOpError("successor operand types must match next block argument types");
+  }
+  return success();
+}
+
+LogicalResult SpatDeferredCommunicationOp::verify() {
+  if (getSources().empty())
+    return emitOpError("requires at least one source");
+  if (failed(verifyRegionArguments(getOperation(), getBody(), getSources(), "spat.deferred_communication")))
+    return failure();
+  return verifyYieldTypes(getOperation(), getBody(), getOperation()->getResultTypes(), "spat.deferred_communication");
+}
+
 template <typename ComputeBatchOpTy>
 LogicalResult verifyComputeBatchLikeOp(ComputeBatchOpTy batch, StringRef opName) {
   int32_t count = batch.getLaneCount();
@@ -808,7 +879,7 @@ LogicalResult verifyComputeBatchLikeOp(ComputeBatchOpTy batch, StringRef opName)
     return failure();
   if (failed(verifyOnlyConstantExternalValues(batch.getOperation(), batch.getBody(), opName)))
     return failure();
-  return verifyBatchBody(batch, block);
+  return verifyBatchBody(batch, block, !isa<SpatScheduledComputeBatch>(batch.getOperation()));
 }
 
 LogicalResult SpatGraphComputeBatch::verify() { return verifyComputeBatchLikeOp(*this, "spat.graph_compute_batch"); }

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

@@ -1,656 +0,0 @@
-#include "CommunicationPlan.hpp"
-
-#include "mlir/IR/Diagnostics.h"
-
-#include "llvm/ADT/DenseSet.h"
-#include "llvm/ADT/STLExtras.h"
-
-#include <algorithm>
-
-using namespace mlir;
-
-namespace onnx_mlir::spatial {
-namespace {
-
-StringRef stringifyExchangeKind(CommunicationExchangeKind kind) {
-  switch (kind) {
-  case CommunicationExchangeKind::HostForward: return "HostForward";
-  case CommunicationExchangeKind::OrdinaryValue: return "OrdinaryValue";
-  case CommunicationExchangeKind::PackedScalarRun: return "PackedScalarRun";
-  case CommunicationExchangeKind::IndexedBatchRun: return "IndexedBatchRun";
-  case CommunicationExchangeKind::ProjectedInput: return "ProjectedInput";
-  }
-  llvm_unreachable("unknown communication exchange kind");
-}
-
-unsigned getExchangeKindOrder(CommunicationExchangeKind kind) {
-  switch (kind) {
-  case CommunicationExchangeKind::HostForward: return 0;
-  case CommunicationExchangeKind::OrdinaryValue: return 1;
-  case CommunicationExchangeKind::PackedScalarRun: return 2;
-  case CommunicationExchangeKind::IndexedBatchRun: return 3;
-  case CommunicationExchangeKind::ProjectedInput: return 4;
-  }
-  llvm_unreachable("unknown communication exchange kind");
-}
-
-void emitPlanError(Operation* anchor, StringRef message) {
-  if (anchor)
-    anchor->emitError(message);
-}
-
-struct PeerPair {
-  int64_t minCore = 0;
-  int64_t maxCore = 0;
-
-  bool operator==(const PeerPair& other) const {
-    return minCore == other.minCore && maxCore == other.maxCore;
-  }
-};
-
-struct PeerPairInfo {
-  static PeerPair getEmptyKey() {
-    return {llvm::DenseMapInfo<int64_t>::getEmptyKey(), llvm::DenseMapInfo<int64_t>::getEmptyKey()};
-  }
-  static PeerPair getTombstoneKey() {
-    return {llvm::DenseMapInfo<int64_t>::getTombstoneKey(), llvm::DenseMapInfo<int64_t>::getTombstoneKey()};
-  }
-  static unsigned getHashValue(PeerPair key) { return llvm::hash_combine(key.minCore, key.maxCore); }
-  static bool isEqual(PeerPair lhs, PeerPair rhs) { return lhs == rhs; }
-};
-
-PeerPair getPeerPair(const CommunicationExchange& exchange) {
-  return {std::min(exchange.sourceCore, exchange.targetCore), std::max(exchange.sourceCore, exchange.targetCore)};
-}
-
-bool exchangeOrderLess(const CommunicationExchange& lhs, const CommunicationExchange& rhs) {
-  if (getExchangeKindOrder(lhs.kind) != getExchangeKindOrder(rhs.kind))
-    return getExchangeKindOrder(lhs.kind) < getExchangeKindOrder(rhs.kind);
-  if (lhs.sourceClass != rhs.sourceClass)
-    return lhs.sourceClass < rhs.sourceClass;
-  if (lhs.targetClass != rhs.targetClass)
-    return lhs.targetClass < rhs.targetClass;
-  if (lhs.producerOrder != rhs.producerOrder)
-    return lhs.producerOrder < rhs.producerOrder;
-  if (lhs.producer.resultIndex != rhs.producer.resultIndex)
-    return lhs.producer.resultIndex < rhs.producer.resultIndex;
-  if (lhs.producer.instance.laneStart != rhs.producer.instance.laneStart)
-    return lhs.producer.instance.laneStart < rhs.producer.instance.laneStart;
-  if (lhs.producer.instance.laneCount != rhs.producer.instance.laneCount)
-    return lhs.producer.instance.laneCount < rhs.producer.instance.laneCount;
-  if (lhs.sourceLane != rhs.sourceLane)
-    return lhs.sourceLane < rhs.sourceLane;
-  if (lhs.targetLane != rhs.targetLane)
-    return lhs.targetLane < rhs.targetLane;
-  if (lhs.ordinal != rhs.ordinal)
-    return lhs.ordinal < rhs.ordinal;
-  return lhs.channelId < rhs.channelId;
-}
-
-void sortExchangeIds(ArrayRef<CommunicationExchange> exchanges,
-                     SmallVectorImpl<CommunicationExchangeId>& ids) {
-  llvm::stable_sort(ids, [&](CommunicationExchangeId lhs, CommunicationExchangeId rhs) {
-    return exchangeOrderLess(exchanges[lhs.value], exchanges[rhs.value]);
-  });
-}
-
-void addEvent(CommunicationPlan& plan,
-              DenseMap<ClassId, SmallVector<CommunicationEvent, 16>>& eventsByClass,
-              CommunicationExchangeId id,
-              CommunicationEventKind kind,
-              ClassId classId,
-              int64_t core,
-              unsigned& order) {
-  eventsByClass[classId].push_back(CommunicationEvent {id, kind, classId, core, order++});
-  (void)plan;
-}
-
-void buildEventStreams(CommunicationPlan& plan,
-                       SmallVectorImpl<CommunicationExchange>& exchanges,
-                       DenseMap<ClassId, SmallVector<CommunicationEvent, 16>>& eventsByClass) {
-  DenseMap<PeerPair, SmallVector<CommunicationExchangeId, 8>, PeerPairInfo> exchangesByPair;
-  for (const CommunicationExchange& exchange : exchanges)
-    exchangesByPair[getPeerPair(exchange)].push_back(exchange.id);
-
-  SmallVector<PeerPair, 8> pairs;
-  pairs.reserve(exchangesByPair.size());
-  for (const auto& entry : exchangesByPair)
-    pairs.push_back(entry.first);
-  llvm::sort(pairs, [](PeerPair lhs, PeerPair rhs) {
-    if (lhs.minCore != rhs.minCore)
-      return lhs.minCore < rhs.minCore;
-    return lhs.maxCore < rhs.maxCore;
-  });
-
-  unsigned order = 0;
-  for (PeerPair pair : pairs) {
-    SmallVector<CommunicationExchangeId, 8> hiToLo;
-    SmallVector<CommunicationExchangeId, 8> loToHi;
-    for (CommunicationExchangeId id : exchangesByPair.lookup(pair)) {
-      const CommunicationExchange& exchange = exchanges[id.value];
-      if (exchange.sourceCore == pair.maxCore && exchange.targetCore == pair.minCore)
-        hiToLo.push_back(id);
-      else
-        loToHi.push_back(id);
-    }
-
-    sortExchangeIds(exchanges, hiToLo);
-    sortExchangeIds(exchanges, loToHi);
-
-    for (CommunicationExchangeId id : hiToLo) {
-      const CommunicationExchange& exchange = exchanges[id.value];
-      addEvent(plan, eventsByClass, id, CommunicationEventKind::Receive, exchange.targetClass, exchange.targetCore, order);
-    }
-    for (CommunicationExchangeId id : hiToLo) {
-      const CommunicationExchange& exchange = exchanges[id.value];
-      addEvent(plan, eventsByClass, id, CommunicationEventKind::Send, exchange.sourceClass, exchange.sourceCore, order);
-    }
-    for (CommunicationExchangeId id : loToHi) {
-      const CommunicationExchange& exchange = exchanges[id.value];
-      addEvent(plan, eventsByClass, id, CommunicationEventKind::Send, exchange.sourceClass, exchange.sourceCore, order);
-    }
-    for (CommunicationExchangeId id : loToHi) {
-      const CommunicationExchange& exchange = exchanges[id.value];
-      addEvent(plan, eventsByClass, id, CommunicationEventKind::Receive, exchange.targetClass, exchange.targetCore, order);
-    }
-  }
-}
-
-} // namespace
-
-FailureOr<CommunicationPlan>
-CommunicationPlan::build(Operation* anchor,
-                         ArrayRef<CommunicationClassInfo> classes,
-                         const DenseMap<CpuId, ClassId>& cpuToClass,
-                         const ProducerSourceClassMap& producerSourceClasses,
-                         const DenseMap<ProducerKey, Type, ProducerKeyInfo>& producerPayloadTypes,
-                         const DenseMap<ProducerKey, SmallVector<ClassId, 4>, ProducerKeyInfo>& producerDestClasses,
-                         const DenseMap<ProducerKey, SmallVector<ClassId, 4>, ProducerKeyInfo>& ordinaryProducerDestClasses,
-                         ProjectedInputPlanMap& projectedPlans,
-                         const DenseMap<Value, ClassId>& hostOutputOwners,
-                         const DenseMap<Value, bool>& liveExternalUseCache,
-                         int64_t& nextChannelId) {
-  CommunicationPlan plan;
-
-  DenseMap<ClassId, CommunicationClassInfo> classById;
-  for (CommunicationClassInfo klass : classes)
-    classById[klass.classId] = std::move(klass);
-
-  auto appendExchange = [&](CommunicationExchangeKind kind,
-                            ProducerKey producer,
-                            ClassId sourceClass,
-                            ClassId targetClass,
-                            Type payloadType,
-                            int64_t sourceCore,
-                            int64_t targetCore,
-                            unsigned sourceLane,
-                            unsigned targetLane,
-                            unsigned ordinal,
-                            Value originalHostOutput = Value(),
-                            const ProjectedInputTransferFragment* fragment = nullptr) -> FailureOr<CommunicationExchangeId> {
-    if (sourceCore == targetCore)
-      return failure();
-
-    auto sourceInfoIt = classById.find(sourceClass);
-    auto targetInfoIt = classById.find(targetClass);
-    if (sourceInfoIt == classById.end() || targetInfoIt == classById.end())
-      return failure();
-
-    CommunicationExchangeId exchangeId {plan.exchanges.size()};
-    plan.exchanges.push_back(CommunicationExchange {
-      exchangeId,
-      kind,
-      PlaceholderId {static_cast<unsigned>(plan.demands.size())},
-      producer,
-      sourceClass,
-      targetClass,
-      sourceInfoIt->second.rank,
-      targetInfoIt->second.rank,
-      payloadType,
-      nextChannelId++,
-      sourceCore,
-      targetCore,
-      sourceLane,
-      targetLane,
-      ordinal,
-      exchangeId.value,
-      originalHostOutput,
-      fragment,
-    });
-    return exchangeId;
-  };
-
-  auto getPayloadType = [&](ProducerKey producer) -> Type {
-    auto typeIt = producerPayloadTypes.find(producer);
-    if (typeIt != producerPayloadTypes.end())
-      return typeIt->second;
-    if (!producer.instance.op || producer.resultIndex >= producer.instance.op->getNumResults())
-      return Type();
-    return producer.instance.op->getResult(producer.resultIndex).getType();
-  };
-
-  for (const auto& entry : producerDestClasses) {
-    ProducerKey producer = entry.first;
-    auto sourceClassIt = producerSourceClasses.find(producer);
-    if (sourceClassIt == producerSourceClasses.end()) {
-      emitPlanError(anchor, "communication plan could not find producer source class");
-      return failure();
-    }
-    ClassId sourceClass = sourceClassIt->second;
-    auto sourceInfoIt = classById.find(sourceClass);
-    if (sourceInfoIt == classById.end()) {
-      emitPlanError(anchor, "communication plan references an unknown producer source class");
-      return failure();
-    }
-    const CommunicationClassInfo& sourceInfo = sourceInfoIt->second;
-    Type payloadType = getPayloadType(producer);
-    if (!payloadType) {
-      emitPlanError(anchor, "communication plan could not determine producer payload type");
-      return failure();
-    }
-
-    for (ClassId targetClass : entry.second) {
-      if (targetClass == sourceClass)
-        continue;
-
-      auto targetInfoIt = classById.find(targetClass);
-      if (targetInfoIt == classById.end()) {
-        emitPlanError(anchor, "communication plan references an unknown target class");
-        return failure();
-      }
-      const CommunicationClassInfo& targetInfo = targetInfoIt->second;
-
-      CommunicationExchangeKind kind = CommunicationExchangeKind::OrdinaryValue;
-      if (sourceInfo.isBatch && !targetInfo.isBatch)
-        kind = CommunicationExchangeKind::PackedScalarRun;
-      else if (sourceInfo.isBatch && targetInfo.isBatch)
-        kind = CommunicationExchangeKind::IndexedBatchRun;
-      else if (!llvm::is_contained(ordinaryProducerDestClasses.lookup(producer), targetClass))
-        continue;
-
-      SmallVector<CommunicationExchangeId, 8>* bucket = nullptr;
-      ProducerTargetKey lookupKey {producer, targetClass};
-      if (kind == CommunicationExchangeKind::OrdinaryValue)
-        bucket = &plan.ordinaryExchangesByProducerTarget[lookupKey];
-      else if (kind == CommunicationExchangeKind::PackedScalarRun)
-        bucket = &plan.packedScalarRunExchangesByProducerTarget[lookupKey];
-      else if (kind == CommunicationExchangeKind::IndexedBatchRun)
-        bucket = &plan.indexedBatchRunExchangesByProducerTarget[lookupKey];
-
-      if (!bucket)
-        continue;
-
-      if (!sourceInfo.isBatch) {
-        if (sourceInfo.cpus.empty()) {
-          emitPlanError(anchor, "communication plan source class has no CPU");
-          return failure();
-        }
-        int64_t sourceCore = static_cast<int64_t>(sourceInfo.cpus.front());
-        unsigned ordinal = 0;
-        for (auto [targetLane, targetCpu] : llvm::enumerate(targetInfo.cpus)) {
-          auto exchangeId = appendExchange(kind,
-                                           producer,
-                                           sourceClass,
-                                           targetClass,
-                                           payloadType,
-                                           sourceCore,
-                                           static_cast<int64_t>(targetCpu),
-                                           0,
-                                           static_cast<unsigned>(targetLane),
-                                           ordinal++);
-          if (failed(exchangeId)) {
-            emitPlanError(anchor, "communication plan failed to append scalar source exchange");
-            return failure();
-          }
-          bucket->push_back(*exchangeId);
-          if (!targetInfo.isBatch)
-            break;
-        }
-        continue;
-      }
-
-      if (sourceInfo.cpus.empty()) {
-        emitPlanError(anchor, "communication plan batch source class has no CPUs");
-        return failure();
-      }
-      if (targetInfo.isBatch && sourceInfo.cpus.size() != targetInfo.cpus.size()) {
-        emitPlanError(anchor, "communication plan cannot map batch classes of different sizes");
-        return failure();
-      }
-
-      for (auto [sourceLane, sourceCpu] : llvm::enumerate(sourceInfo.cpus)) {
-        CpuId targetCpu = targetInfo.isBatch ? targetInfo.cpus[sourceLane] : targetInfo.cpus.front();
-        auto exchangeId = appendExchange(kind,
-                                         producer,
-                                         sourceClass,
-                                         targetClass,
-                                         payloadType,
-                                         static_cast<int64_t>(sourceCpu),
-                                         static_cast<int64_t>(targetCpu),
-                                         static_cast<unsigned>(sourceLane),
-                                         targetInfo.isBatch ? static_cast<unsigned>(sourceLane) : 0,
-                                         static_cast<unsigned>(bucket->size()));
-        if (failed(exchangeId)) {
-          emitPlanError(anchor, "communication plan failed to append batch source exchange");
-          return failure();
-        }
-        bucket->push_back(*exchangeId);
-      }
-    }
-  }
-
-  for (const auto& entry : producerSourceClasses) {
-    ProducerKey producer = entry.first;
-    ClassId sourceClass = entry.second;
-    if (!producer.instance.op || producer.resultIndex >= producer.instance.op->getNumResults())
-      continue;
-
-    Value originalOutput = producer.instance.op->getResult(producer.resultIndex);
-    auto liveIt = liveExternalUseCache.find(originalOutput);
-    if (liveIt == liveExternalUseCache.end() || !liveIt->second)
-      continue;
-
-    auto ownerIt = hostOutputOwners.find(originalOutput);
-    if (ownerIt == hostOutputOwners.end())
-      continue;
-    ClassId ownerClass = ownerIt->second;
-    if (ownerClass == sourceClass)
-      continue;
-
-    auto sourceInfoIt = classById.find(sourceClass);
-    auto ownerInfoIt = classById.find(ownerClass);
-    if (sourceInfoIt == classById.end() || ownerInfoIt == classById.end()) {
-      emitPlanError(anchor, "communication plan references an unknown host-forward class");
-      return failure();
-    }
-    const CommunicationClassInfo& sourceInfo = sourceInfoIt->second;
-    const CommunicationClassInfo& ownerInfo = ownerInfoIt->second;
-    if (sourceInfo.isBatch != ownerInfo.isBatch) {
-      emitPlanError(anchor, "communication plan does not support mixed scalar/batch host forwarding");
-      return failure();
-    }
-    if (sourceInfo.isBatch && sourceInfo.cpus.size() != ownerInfo.cpus.size()) {
-      emitPlanError(anchor, "communication plan cannot host-forward between batch classes of different sizes");
-      return failure();
-    }
-
-    Type payloadType = getPayloadType(producer);
-    if (!payloadType) {
-      emitPlanError(anchor, "communication plan could not determine host-forward payload type");
-      return failure();
-    }
-
-    HostForwardKey lookupKey {originalOutput, sourceClass, ownerClass};
-    SmallVector<CommunicationExchangeId, 8>& bucket = plan.hostForwardExchangesByOutput[lookupKey];
-    if (!sourceInfo.isBatch) {
-      if (!bucket.empty())
-        continue;
-      auto exchangeId = appendExchange(CommunicationExchangeKind::HostForward,
-                                       producer,
-                                       sourceClass,
-                                       ownerClass,
-                                       payloadType,
-                                       static_cast<int64_t>(sourceInfo.cpus.front()),
-                                       static_cast<int64_t>(ownerInfo.cpus.front()),
-                                       0,
-                                       0,
-                                       0,
-                                       originalOutput);
-      if (failed(exchangeId)) {
-        emitPlanError(anchor, "communication plan failed to append scalar host-forward exchange");
-        return failure();
-      }
-      bucket.push_back(*exchangeId);
-      continue;
-    }
-
-    if (bucket.size() >= sourceInfo.cpus.size())
-      continue;
-    unsigned ordinal = static_cast<unsigned>(bucket.size());
-    if (producer.instance.laneStart >= sourceInfo.cpus.size()) {
-      emitPlanError(anchor, "communication plan host-forward producer lane is out of range");
-      return failure();
-    }
-    unsigned lane = producer.instance.laneStart;
-    auto exchangeId = appendExchange(CommunicationExchangeKind::HostForward,
-                                     producer,
-                                     sourceClass,
-                                     ownerClass,
-                                     payloadType,
-                                     static_cast<int64_t>(sourceInfo.cpus[lane]),
-                                     static_cast<int64_t>(ownerInfo.cpus[lane]),
-                                     lane,
-                                     lane,
-                                     ordinal,
-                                     originalOutput);
-    if (failed(exchangeId)) {
-      emitPlanError(anchor, "communication plan failed to append batch host-forward exchange");
-      return failure();
-    }
-    bucket.push_back(*exchangeId);
-  }
-
-  DenseSet<const ProjectedInputTransferFragment*> seenFragments;
-  for (auto& extractEntry : projectedPlans) {
-    for (auto& classEntry : extractEntry.second) {
-      ClassId targetClass = classEntry.first;
-      ProjectedInputTransferPlan& inputPlan = classEntry.second;
-
-      auto targetInfoIt = classById.find(targetClass);
-      if (targetInfoIt == classById.end()) {
-        emitPlanError(anchor, "communication plan references an unknown target class");
-        return failure();
-      }
-      const CommunicationClassInfo& targetInfo = targetInfoIt->second;
-
-      for (ProjectedInputTransferFragment& fragment : inputPlan.fragments) {
-        if (!seenFragments.insert(&fragment).second) {
-          emitPlanError(anchor, "communication plan found the same projected fragment twice");
-          return failure();
-        }
-
-        auto sourceClassIt = cpuToClass.find(fragment.sourceCoreId);
-        if (sourceClassIt == cpuToClass.end()) {
-          emitPlanError(anchor, "communication plan could not map projected source core to a class");
-          return failure();
-        }
-
-        ClassId sourceClass = sourceClassIt->second;
-        auto sourceInfoIt = classById.find(sourceClass);
-        if (sourceInfoIt == classById.end()) {
-          emitPlanError(anchor, "communication plan references an unknown source class");
-          return failure();
-        }
-        const CommunicationClassInfo& sourceInfo = sourceInfoIt->second;
-
-        plan.demands.push_back(InputDemand {{static_cast<unsigned>(plan.demands.size())},
-                                            fragment.producer,
-                                            targetClass,
-                                            inputPlan.layout.fragmentType});
-        plan.publications.push_back(ProducerPublication {fragment.producer, sourceClass, inputPlan.layout.fragmentType});
-
-        if (fragment.sourceCoreId == fragment.targetCoreId)
-          continue;
-
-        if (sourceInfo.rank == targetInfo.rank && sourceClass != targetClass) {
-          emitPlanError(anchor, "communication plan found a same-rank inter-class projected exchange");
-          return failure();
-        }
-
-        FailureOr<CommunicationExchangeId> exchangeId = appendExchange(
-          CommunicationExchangeKind::ProjectedInput,
-          fragment.producer,
-          sourceClass,
-          targetClass,
-          inputPlan.layout.fragmentType,
-          fragment.sourceCoreId,
-          fragment.targetCoreId,
-          0,
-          fragment.targetLane,
-          fragment.ordinal,
-          Value(),
-          &fragment);
-        if (failed(exchangeId)) {
-          emitPlanError(anchor, "communication plan failed to append projected input exchange");
-          return failure();
-        }
-        fragment.channelId = plan.exchanges[exchangeId->value].channelId;
-        plan.projectedExchangesByFragment[&fragment].push_back(*exchangeId);
-      }
-    }
-  }
-
-  buildEventStreams(plan, plan.exchanges, plan.eventsByClass);
-  if (failed(plan.verify(anchor)))
-    return failure();
-  return plan;
-}
-
-LogicalResult CommunicationPlan::verify(Operation* anchor) const {
-  DenseSet<int64_t> channelIds;
-  SmallVector<const CommunicationExchange*, 8> malformed;
-  size_t malformedCount = 0;
-
-  auto recordMalformed = [&](const CommunicationExchange& exchange) {
-    ++malformedCount;
-    if (malformed.size() < 8)
-      malformed.push_back(&exchange);
-  };
-
-  for (auto [index, exchange] : llvm::enumerate(exchanges)) {
-    bool bad = false;
-    if (exchange.id.value != index)
-      bad = true;
-    if (exchange.sourceCore == exchange.targetCore)
-      bad = true;
-    if (!exchange.payloadType)
-      bad = true;
-    if (!channelIds.insert(exchange.channelId).second)
-      bad = true;
-    if (bad)
-      recordMalformed(exchange);
-  }
-
-  DenseMap<size_t, unsigned> sendCounts;
-  DenseMap<size_t, unsigned> receiveCounts;
-  for (const auto& classEvents : eventsByClass) {
-    std::optional<unsigned> previousOrder;
-    DenseSet<unsigned> classOrders;
-    for (const CommunicationEvent& event : classEvents.second) {
-      if (event.exchangeId.value >= exchanges.size()) {
-        if (anchor)
-          anchor->emitError("communication plan event references an invalid exchange id");
-        return failure();
-      }
-
-      const CommunicationExchange& exchange = exchanges[event.exchangeId.value];
-      if (previousOrder && event.order < *previousOrder) {
-        if (anchor)
-          anchor->emitError("communication plan event stream is not sorted");
-        return failure();
-      }
-      previousOrder = event.order;
-      if (!classOrders.insert(event.order).second) {
-        if (anchor)
-          anchor->emitError("communication plan event stream contains duplicate event order");
-        return failure();
-      }
-
-      if (event.kind == CommunicationEventKind::Send) {
-        ++sendCounts[event.exchangeId.value];
-        if (event.classId != exchange.sourceClass || event.core != exchange.sourceCore)
-          recordMalformed(exchange);
-      }
-      else {
-        ++receiveCounts[event.exchangeId.value];
-        if (event.classId != exchange.targetClass || event.core != exchange.targetCore)
-          recordMalformed(exchange);
-      }
-    }
-  }
-
-  for (const CommunicationExchange& exchange : exchanges) {
-    if (sendCounts[exchange.id.value] != 1 || receiveCounts[exchange.id.value] != 1)
-      recordMalformed(exchange);
-  }
-
-  auto verifyLookup = [&](ArrayRef<CommunicationExchangeId> ids) -> LogicalResult {
-    for (CommunicationExchangeId id : ids)
-      if (id.value >= exchanges.size())
-        return failure();
-    return success();
-  };
-
-  for (const auto& entry : projectedExchangesByFragment)
-    if (failed(verifyLookup(entry.second))) {
-      if (anchor)
-        anchor->emitError("communication plan projected lookup references an invalid exchange id");
-      return failure();
-    }
-
-  if (malformedCount != 0) {
-    if (anchor) {
-      InFlightDiagnostic diag = anchor->emitError("communication plan contains malformed exchanges");
-      for (const CommunicationExchange* exchange : malformed) {
-        diag << " [exchange=" << exchange->id.value << " kind=" << stringifyExchangeKind(exchange->kind)
-             << " sourceClass=" << exchange->sourceClass << " sourceCore=" << exchange->sourceCore
-             << " targetClass=" << exchange->targetClass << " targetCore=" << exchange->targetCore
-             << " channelId=" << exchange->channelId << " payloadType=" << exchange->payloadType << "]";
-      }
-      if (malformedCount > malformed.size())
-        diag << " additionalMalformed=" << (malformedCount - malformed.size());
-    }
-    return failure();
-  }
-
-  return success();
-}
-
-ArrayRef<CommunicationEvent> CommunicationPlan::getEventsForClass(ClassId classId) const {
-  auto it = eventsByClass.find(classId);
-  if (it == eventsByClass.end())
-    return {};
-  return it->second;
-}
-
-ArrayRef<CommunicationExchangeId>
-CommunicationPlan::getOrdinaryValueExchanges(ProducerKey producer, ClassId targetClass) const {
-  auto it = ordinaryExchangesByProducerTarget.find({producer, targetClass});
-  if (it == ordinaryExchangesByProducerTarget.end())
-    return {};
-  return it->second;
-}
-
-ArrayRef<CommunicationExchangeId>
-CommunicationPlan::getProjectedInputExchanges(const ProjectedInputTransferFragment& fragment) const {
-  auto it = projectedExchangesByFragment.find(&fragment);
-  if (it == projectedExchangesByFragment.end())
-    return {};
-  return it->second;
-}
-
-ArrayRef<CommunicationExchangeId>
-CommunicationPlan::getPackedScalarRunExchanges(ProducerKey producer, ClassId targetClass) const {
-  auto it = packedScalarRunExchangesByProducerTarget.find({producer, targetClass});
-  if (it == packedScalarRunExchangesByProducerTarget.end())
-    return {};
-  return it->second;
-}
-
-ArrayRef<CommunicationExchangeId>
-CommunicationPlan::getIndexedBatchRunExchanges(ProducerKey producer, ClassId targetClass) const {
-  auto it = indexedBatchRunExchangesByProducerTarget.find({producer, targetClass});
-  if (it == indexedBatchRunExchangesByProducerTarget.end())
-    return {};
-  return it->second;
-}
-
-ArrayRef<CommunicationExchangeId>
-CommunicationPlan::getHostForwardExchanges(Value originalOutput, ClassId sourceClass, ClassId ownerClass) const {
-  auto it = hostForwardExchangesByOutput.find({originalOutput, sourceClass, ownerClass});
-  if (it == hostForwardExchangesByOutput.end())
-    return {};
-  return it->second;
-}
-
-} // namespace onnx_mlir::spatial

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/CommunicationPlan.hpp

@@ -1,229 +0,0 @@
-#pragma once
-
-#include "mlir/IR/Operation.h"
-#include "mlir/IR/Types.h"
-#include "mlir/Support/LLVM.h"
-
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/DenseSet.h"
-#include "llvm/ADT/SmallVector.h"
-
-#include <cstddef>
-#include <cstdint>
-#include <optional>
-
-#include "MergeScheduleKeys.hpp"
-#include "ProjectedFragments.hpp"
-
-namespace onnx_mlir::spatial {
-
-enum class TransferKind {
-  DirectValue,
-  LocalFragment,
-  SameClassForward,
-  RemoteChannel,
-  WholeTensorBarrier
-};
-
-enum class CommunicationExchangeKind {
-  OrdinaryValue,
-  ProjectedInput,
-  PackedScalarRun,
-  IndexedBatchRun,
-  HostForward
-};
-
-enum class CommunicationEventKind {
-  Send,
-  Receive
-};
-
-struct CommunicationExchangeId {
-  size_t value = 0;
-
-  bool operator==(const CommunicationExchangeId& other) const { return value == other.value; }
-  operator size_t() const { return value; }
-};
-
-struct CommunicationExchangeIdInfo {
-  static CommunicationExchangeId getEmptyKey() { return {llvm::DenseMapInfo<size_t>::getEmptyKey()}; }
-  static CommunicationExchangeId getTombstoneKey() { return {llvm::DenseMapInfo<size_t>::getTombstoneKey()}; }
-  static unsigned getHashValue(CommunicationExchangeId id) { return llvm::hash_value(id.value); }
-  static bool isEqual(CommunicationExchangeId lhs, CommunicationExchangeId rhs) { return lhs == rhs; }
-};
-
-struct PlaceholderId {
-  unsigned value = 0;
-
-  bool operator==(const PlaceholderId& other) const { return value == other.value; }
-};
-
-struct PlaceholderIdInfo {
-  static PlaceholderId getEmptyKey() { return {static_cast<unsigned>(-1)}; }
-  static PlaceholderId getTombstoneKey() { return {static_cast<unsigned>(-2)}; }
-  static unsigned getHashValue(PlaceholderId id) { return llvm::hash_value(id.value); }
-  static bool isEqual(PlaceholderId lhs, PlaceholderId rhs) { return lhs == rhs; }
-};
-
-struct InputDemand {
-  PlaceholderId placeholder;
-  ProducerKey producer;
-  ClassId targetClass = 0;
-  mlir::Type requiredType;
-};
-
-struct ProducerPublication {
-  ProducerKey producer;
-  ClassId sourceClass = 0;
-  mlir::Type payloadType;
-};
-
-struct CommunicationClassInfo {
-  ClassId classId = 0;
-  size_t rank = 0;
-  bool isBatch = false;
-  mlir::Operation* op = nullptr;
-  llvm::SmallVector<CpuId, 8> cpus;
-};
-
-struct CommunicationExchange {
-  CommunicationExchangeId id;
-  CommunicationExchangeKind kind = CommunicationExchangeKind::OrdinaryValue;
-
-  PlaceholderId placeholder;
-  ProducerKey producer;
-  ClassId sourceClass = 0;
-  ClassId targetClass = 0;
-
-  size_t sourceRank = 0;
-  size_t targetRank = 0;
-  mlir::Type payloadType;
-  int64_t channelId = 0;
-  int64_t sourceCore = 0;
-  int64_t targetCore = 0;
-
-  unsigned sourceLane = 0;
-  unsigned targetLane = 0;
-  unsigned ordinal = 0;
-  uint64_t producerOrder = 0;
-
-  mlir::Value originalHostOutput;
-  const ProjectedInputTransferFragment* fragment = nullptr;
-};
-
-using ExchangeDescriptor = CommunicationExchange;
-
-struct CommunicationEvent {
-  CommunicationExchangeId exchangeId;
-  CommunicationEventKind kind = CommunicationEventKind::Send;
-  ClassId classId = 0;
-  int64_t core = 0;
-  unsigned order = 0;
-};
-
-class CommunicationPlan {
-public:
-  using ProjectedInputPlanMap =
-    llvm::DenseMap<mlir::Operation*, llvm::DenseMap<ClassId, ProjectedInputTransferPlan>>;
-  using ProducerSourceClassMap = llvm::DenseMap<ProducerKey, ClassId, ProducerKeyInfo>;
-
-  static mlir::FailureOr<CommunicationPlan>
-  build(mlir::Operation* anchor,
-        llvm::ArrayRef<CommunicationClassInfo> classes,
-        const llvm::DenseMap<CpuId, ClassId>& cpuToClass,
-        const ProducerSourceClassMap& producerSourceClasses,
-        const llvm::DenseMap<ProducerKey, mlir::Type, ProducerKeyInfo>& producerPayloadTypes,
-        const llvm::DenseMap<ProducerKey, llvm::SmallVector<ClassId, 4>, ProducerKeyInfo>& producerDestClasses,
-        const llvm::DenseMap<ProducerKey, llvm::SmallVector<ClassId, 4>, ProducerKeyInfo>& ordinaryProducerDestClasses,
-        ProjectedInputPlanMap& projectedPlans,
-        const llvm::DenseMap<mlir::Value, ClassId>& hostOutputOwners,
-        const llvm::DenseMap<mlir::Value, bool>& liveExternalUseCache,
-        int64_t& nextChannelId);
-
-  mlir::LogicalResult verify(mlir::Operation* anchor) const;
-
-  llvm::ArrayRef<CommunicationEvent> getEventsForClass(ClassId classId) const;
-  llvm::ArrayRef<CommunicationExchange> getExchanges() const { return exchanges; }
-  const CommunicationExchange& getExchange(CommunicationExchangeId id) const { return exchanges[id.value]; }
-  const CommunicationExchange& getExchange(size_t id) const { return exchanges[id]; }
-
-  llvm::ArrayRef<CommunicationExchangeId>
-  getOrdinaryValueExchanges(ProducerKey producer, ClassId targetClass) const;
-  llvm::ArrayRef<CommunicationExchangeId>
-  getProjectedInputExchanges(const ProjectedInputTransferFragment& fragment) const;
-  llvm::ArrayRef<CommunicationExchangeId>
-  getPackedScalarRunExchanges(ProducerKey producer, ClassId targetClass) const;
-  llvm::ArrayRef<CommunicationExchangeId>
-  getIndexedBatchRunExchanges(ProducerKey producer, ClassId targetClass) const;
-  llvm::ArrayRef<CommunicationExchangeId>
-  getHostForwardExchanges(mlir::Value originalOutput, ClassId sourceClass, ClassId ownerClass) const;
-
-private:
-  struct ProducerTargetKey {
-    ProducerKey producer;
-    ClassId targetClass = 0;
-
-    bool operator==(const ProducerTargetKey& other) const {
-      return producer == other.producer && targetClass == other.targetClass;
-    }
-  };
-
-  struct ProducerTargetKeyInfo {
-    static ProducerTargetKey getEmptyKey() {
-      return {ProducerKeyInfo::getEmptyKey(), llvm::DenseMapInfo<ClassId>::getEmptyKey()};
-    }
-    static ProducerTargetKey getTombstoneKey() {
-      return {ProducerKeyInfo::getTombstoneKey(), llvm::DenseMapInfo<ClassId>::getTombstoneKey()};
-    }
-    static unsigned getHashValue(const ProducerTargetKey& key) {
-      return llvm::hash_combine(ProducerKeyInfo::getHashValue(key.producer), key.targetClass);
-    }
-    static bool isEqual(const ProducerTargetKey& lhs, const ProducerTargetKey& rhs) { return lhs == rhs; }
-  };
-
-  struct HostForwardKey {
-    mlir::Value originalOutput;
-    ClassId sourceClass = 0;
-    ClassId ownerClass = 0;
-
-    bool operator==(const HostForwardKey& other) const {
-      return originalOutput == other.originalOutput && sourceClass == other.sourceClass
-          && ownerClass == other.ownerClass;
-    }
-  };
-
-  struct HostForwardKeyInfo {
-    static HostForwardKey getEmptyKey() {
-      return {llvm::DenseMapInfo<mlir::Value>::getEmptyKey(),
-              llvm::DenseMapInfo<ClassId>::getEmptyKey(),
-              llvm::DenseMapInfo<ClassId>::getEmptyKey()};
-    }
-    static HostForwardKey getTombstoneKey() {
-      return {llvm::DenseMapInfo<mlir::Value>::getTombstoneKey(),
-              llvm::DenseMapInfo<ClassId>::getTombstoneKey(),
-              llvm::DenseMapInfo<ClassId>::getTombstoneKey()};
-    }
-    static unsigned getHashValue(const HostForwardKey& key) {
-      return llvm::hash_combine(key.originalOutput, key.sourceClass, key.ownerClass);
-    }
-    static bool isEqual(const HostForwardKey& lhs, const HostForwardKey& rhs) { return lhs == rhs; }
-  };
-
-  llvm::SmallVector<InputDemand, 16> demands;
-  llvm::SmallVector<ProducerPublication, 16> publications;
-  llvm::SmallVector<CommunicationExchange, 16> exchanges;
-  llvm::DenseMap<ClassId, llvm::SmallVector<CommunicationEvent, 16>> eventsByClass;
-
-  llvm::DenseMap<const ProjectedInputTransferFragment*, llvm::SmallVector<CommunicationExchangeId, 1>>
-    projectedExchangesByFragment;
-  llvm::DenseMap<ProducerTargetKey, llvm::SmallVector<CommunicationExchangeId, 8>, ProducerTargetKeyInfo>
-    ordinaryExchangesByProducerTarget;
-  llvm::DenseMap<ProducerTargetKey, llvm::SmallVector<CommunicationExchangeId, 8>, ProducerTargetKeyInfo>
-    packedScalarRunExchangesByProducerTarget;
-  llvm::DenseMap<ProducerTargetKey, llvm::SmallVector<CommunicationExchangeId, 8>, ProducerTargetKeyInfo>
-    indexedBatchRunExchangesByProducerTarget;
-  llvm::DenseMap<HostForwardKey, llvm::SmallVector<CommunicationExchangeId, 8>, HostForwardKeyInfo>
-    hostForwardExchangesByOutput;
-};
-
-} // namespace onnx_mlir::spatial

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

@@ -1,134 +0,0 @@
-#include "HostOutputFinalization.hpp"
-
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/IR/BuiltinTypes.h"
-#include "mlir/IR/PatternMatch.h"
-
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/DenseSet.h"
-#include "llvm/ADT/STLExtras.h"
-
-#include "MaterializedClassState.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-
-using namespace mlir;
-
-namespace onnx_mlir::spatial {
-
-LogicalResult finalizeProjectedHostOutputFragments(MaterializerState& state) {
-  if (state.pendingProjectedHostOutputFragments.empty())
-    return success();
-
-  DenseMap<Value, SmallVector<PendingProjectedHostOutputFragment*, 16>> byOutput;
-  for (PendingProjectedHostOutputFragment& fragment : state.pendingProjectedHostOutputFragments)
-    byOutput[fragment.originalOutput].push_back(&fragment);
-
-  SmallVector<Value, 8> outputs;
-  outputs.reserve(byOutput.size());
-
-  auto returnOp = dyn_cast<func::ReturnOp>(state.func.getBody().front().getTerminator());
-  if (!returnOp)
-    return state.func.emitError("expected func.return terminator while finalizing projected host output fragments");
-
-  DenseSet<Value> seenOutputs;
-  for (Value returned : returnOp.getOperands()) {
-    if (!byOutput.contains(returned) || !seenOutputs.insert(returned).second)
-      continue;
-    outputs.push_back(returned);
-  }
-  if (outputs.size() != byOutput.size())
-    return state.func.emitError("projected host output fragments must be keyed by returned logical host outputs");
-
-  for (Value originalOutput : outputs) {
-    if (isa_and_present<SpatScheduledCompute, SpatScheduledComputeBatch>(originalOutput.getDefiningOp())) {
-      return state.func.emitError(
-        "projected host output assembly must be keyed by the original logical host output, not by a materialized scheduled result");
-    }
-
-    auto resultType = dyn_cast<RankedTensorType>(originalOutput.getType());
-    if (!resultType || !resultType.hasStaticShape())
-      return state.func.emitError("projected host output must have static ranked tensor type");
-
-    SmallVector<PendingProjectedHostOutputFragment*, 16>& fragments = byOutput[originalOutput];
-    llvm::sort(fragments, [](const PendingProjectedHostOutputFragment* lhs,
-                             const PendingProjectedHostOutputFragment* rhs) {
-      if (lhs->sourceClass != rhs->sourceClass)
-        return lhs->sourceClass < rhs->sourceClass;
-      if (lhs->publicationResultIndex != rhs->publicationResultIndex)
-        return lhs->publicationResultIndex < rhs->publicationResultIndex;
-      if (lhs->sourceFragmentOrdinal != rhs->sourceFragmentOrdinal)
-        return lhs->sourceFragmentOrdinal < rhs->sourceFragmentOrdinal;
-      return std::lexicographical_compare(lhs->offsets.begin(),
-                                          lhs->offsets.end(),
-                                          rhs->offsets.begin(),
-                                          rhs->offsets.end());
-    });
-
-    state.rewriter.setInsertionPoint(returnOp);
-    Location loc = fragments.front()->loc;
-    SmallVector<Value, 16> blueprintOperands;
-    SmallVector<int64_t, 16> fragmentOperandIndices;
-    SmallVector<int64_t, 16> fragmentSourceOffsets;
-    SmallVector<int64_t, 64> flatOffsets;
-    SmallVector<int64_t, 64> flatSizes;
-    SmallVector<int64_t, 64> flatStrides;
-    DenseMap<Value, int64_t> operandIndicesByValue;
-
-    for (PendingProjectedHostOutputFragment* fragmentRecord : fragments) {
-      if (fragmentRecord->sourceClass >= state.classes.size())
-        return state.func.emitError("projected host output fragment references an invalid source class");
-
-      MaterializedClass& sourceClass = state.classes[fragmentRecord->sourceClass];
-      if (fragmentRecord->publicationResultIndex >= sourceClass.op->getNumResults()) {
-        return sourceClass.op->emitError("projected host output fragment references an invalid publication result")
-               << " sourceClass=" << sourceClass.id
-               << " resultIndex=" << fragmentRecord->publicationResultIndex
-               << " resultCount=" << sourceClass.op->getNumResults();
-      }
-
-      Value operand = sourceClass.op->getResult(fragmentRecord->publicationResultIndex);
-      auto [operandIt, inserted] =
-        operandIndicesByValue.try_emplace(operand, static_cast<int64_t>(blueprintOperands.size()));
-      if (inserted)
-        blueprintOperands.push_back(operand);
-      fragmentOperandIndices.push_back(operandIt->second);
-      fragmentSourceOffsets.push_back(fragmentRecord->sourceElementOffset);
-      llvm::append_range(flatOffsets, fragmentRecord->offsets);
-      llvm::append_range(flatSizes, fragmentRecord->sizes);
-      llvm::append_range(flatStrides, fragmentRecord->strides);
-
-      auto operandType = dyn_cast<RankedTensorType>(operand.getType());
-      if (!operandType || !operandType.hasStaticShape())
-        return state.func.emitError("projected host output assembly requires static ranked tensor operands");
-    }
-
-    if (blueprintOperands.empty())
-      return state.func.emitError("missing projected host output fragments");
-
-    Value input = blueprintOperands.front();
-    ValueRange extraFragments = ValueRange(blueprintOperands).drop_front();
-    auto blueprint = SpatBlueprintOp::create(
-      state.rewriter,
-      loc,
-      resultType,
-      input,
-      extraFragments,
-      state.rewriter.getStringAttr("nchw"),
-      state.rewriter.getStringAttr("fragmented"),
-      state.rewriter.getDenseI64ArrayAttr(flatOffsets),
-      state.rewriter.getDenseI64ArrayAttr(flatSizes),
-      state.rewriter.getStringAttr("identity"),
-      state.rewriter.getStringAttr("fragment_assembly"),
-      state.rewriter.getDenseI64ArrayAttr(fragmentOperandIndices),
-      state.rewriter.getDenseI64ArrayAttr(fragmentSourceOffsets),
-      state.rewriter.getDenseI64ArrayAttr(flatStrides),
-      state.rewriter.getStringAttr("disjoint"),
-      state.rewriter.getStringAttr("complete"));
-
-    state.hostReplacements[originalOutput] = blueprint.getOutput();
-  }
-
-  return success();
-}
-
-} // namespace onnx_mlir::spatial

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/HostOutputFinalization.hpp

@@ -1,11 +0,0 @@
-#pragma once
-
-#include "mlir/Support/LogicalResult.h"
-
-namespace onnx_mlir::spatial {
-
-struct MaterializerState;
-
-mlir::LogicalResult finalizeProjectedHostOutputFragments(MaterializerState& state);
-
-} // namespace onnx_mlir::spatial

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

@@ -1,17 +0,0 @@
-#pragma once
-
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/Support/LogicalResult.h"
-
-#include "Scheduling/MergeSchedule.hpp"
-
-namespace onnx_mlir {
-namespace spatial {
-
-class MergeScheduleMaterializer {
-public:
-  mlir::LogicalResult run(mlir::func::FuncOp func, const MergeScheduleResult& schedule, int64_t& nextChannelId);
-};
-
-} // namespace spatial
-} // namespace onnx_mlir

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializedClassState.hpp

@@ -1,296 +0,0 @@
-#pragma once
-
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/IR/PatternMatch.h"
-#include "mlir/Transforms/FoldUtils.h"
-
-#include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/DenseSet.h"
-#include "llvm/ADT/SmallVector.h"
-
-#include <optional>
-
-#include "MaterializeMergeSchedule.hpp"
-#include "CommunicationPlan.hpp"
-#include "MergeMessages.hpp"
-#include "MergeScheduleKeys.hpp"
-#include "ProjectedFragments.hpp"
-#include "Scheduling/ComputeInstanceUtils.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-
-namespace onnx_mlir::spatial {
-
-struct MaterializedClass {
-  ClassId id = 0;
-  llvm::SmallVector<CpuId, 8> cpus;
-  mlir::Operation* op = nullptr;
-  mlir::Block* body = nullptr;
-  bool isBatch = false;
-
-  llvm::DenseMap<CpuId, unsigned> cpuToLane;
-  llvm::SmallVector<mlir::Value, 8> weights;
-  llvm::SmallVector<mlir::Value, 8> inputs;
-  llvm::SmallVector<mlir::Value, 4> hostOutputs;
-  llvm::DenseMap<mlir::Value, unsigned> publicationOutputToResultIndex;
-  llvm::DenseMap<mlir::Value, mlir::BlockArgument> weightArgs;
-  llvm::DenseMap<mlir::Value, mlir::BlockArgument> inputArgs;
-  llvm::DenseMap<mlir::Value, unsigned> hostOutputToResultIndex;
-};
-
-struct PackedScalarRunSlot {
-  llvm::SmallVector<ProducerKey, 8> keys;
-};
-
-enum class PackedScalarRunKind {
-  Materialized,
-  DeferredReceive,
-  DeferredLocalCompute
-};
-
-struct PackedScalarRunValue {
-  ClassId targetClass = 0;
-  mlir::Operation* sourceOp = nullptr;
-  size_t resultIndex = 0;
-  PackedScalarRunKind kind = PackedScalarRunKind::Materialized;
-
-  mlir::Value packed;
-
-  mlir::RankedTensorType fragmentType;
-  llvm::SmallVector<PackedScalarRunSlot, 8> slots;
-  llvm::SmallVector<CommunicationExchangeId, 8> exchangeIds;
-};
-
-struct IndexedBatchRunValue {
-  ClassId targetClass = 0;
-  mlir::Operation* sourceOp = nullptr;
-  size_t resultIndex = 0;
-  mlir::Value packed;
-  mlir::RankedTensorType fragmentType;
-  llvm::SmallVector<PackedScalarRunSlot, 8> slots;
-  llvm::SmallVector<CommunicationExchangeId, 8> exchangeIds;
-};
-
-struct LogicalSlotRange {
-  SlotId start = 0;
-  SlotId count = 0;
-};
-
-struct MaterializationRunSlot {
-  llvm::SmallVector<ComputeInstance, 8> peers;
-};
-
-using MaterializationRun = llvm::SmallVector<MaterializationRunSlot, 8>;
-
-struct OutputDestinationGroup {
-  llvm::SmallVector<size_t, 4> resultIndices;
-  llvm::SmallVector<ClassId, 4> destinationClasses;
-};
-
-struct BatchRunSendPlan {
-  size_t resultIndex = 0;
-  ClassId destinationClass = 0;
-  llvm::SmallVector<CommunicationExchangeId, 8> exchangeIds;
-};
-
-enum class TensorDemandActionKind {
-  DestinationFanout,
-  SameClassIndexedFragment,
-  ProjectedInputPublication,
-  TerminalBlueprintPublication,
-  WholeTensorBarrier
-};
-
-enum class WholeTensorBarrierReason {
-  FunctionReturnWithoutBlueprint,
-  DenseLogicalConsumer
-};
-
-struct TensorDemandAction {
-  TensorDemandActionKind kind = TensorDemandActionKind::DestinationFanout;
-  std::optional<ClassId> destinationClass;
-  std::optional<WholeTensorBarrierReason> barrierReason;
-};
-
-struct RunOutputDemand {
-  size_t resultIndex = 0;
-  mlir::Value originalOutput;
-  mlir::RankedTensorType fragmentType;
-  llvm::SmallVector<TensorDemandAction, 4> actions;
-};
-
-struct CompactRunPlan {
-  llvm::SmallVector<RunOutputDemand, 4> outputs;
-};
-
-enum class BatchInputDemandKind {
-  LaneFragment,
-  ProjectedFragment,
-  WholeTensorBarrier
-};
-
-struct BatchInputDemand {
-  BatchInputDemandKind kind = BatchInputDemandKind::LaneFragment;
-  std::optional<ProducerKey> wholeTensorProducer;
-};
-
-struct CommunicationExchangeEmissionState {
-  bool producerBound = false;
-  bool sendEmitted = false;
-  bool receiveEmitted = false;
-
-  mlir::Value producerPayload;
-  mlir::Value receiveValue;
-
-  mlir::Operation* producerAnchor = nullptr;
-  mlir::Operation* sendAnchor = nullptr;
-  mlir::Operation* receiveAnchor = nullptr;
-  std::optional<mlir::Location> loc;
-};
-
-struct CloneIndexingContext {
-  std::optional<mlir::Value> runSlotIndex;
-  std::optional<mlir::Value> projectionSlotIndex;
-};
-
-struct MaterializerState;
-
-class AvailableValueStore {
-public:
-  struct ExactBatchFragmentRecord {
-    ProducerKey key;
-    mlir::Value value;
-  };
-
-  void record(ProducerKey key, ClassId classId, mlir::Value value) {
-    exactValues[key][classId] = value;
-
-    auto batch = mlir::dyn_cast_or_null<SpatComputeBatch>(key.instance.op);
-    if (!batch || key.instance.laneCount == 0)
-      return;
-
-    WholeBatchAssemblyLookupKey lookupKey {batch.getOperation(), key.resultIndex, classId};
-    llvm::SmallVector<ExactBatchFragmentRecord, 16>& bucket = exactBatchFragmentsByProducerResultClass[lookupKey];
-    for (ExactBatchFragmentRecord& record : bucket) {
-      if (!(record.key == key))
-        continue;
-      record.value = value;
-      return;
-    }
-    bucket.push_back({key, value});
-  }
-
-  void recordPackedRun(PackedScalarRunValue run) {
-    size_t runIndex = packedScalarRuns.size();
-    packedScalarRuns.push_back(std::move(run));
-    const PackedScalarRunValue& storedRun = packedScalarRuns[runIndex];
-    WholeBatchAssemblyLookupKey lookupKey {storedRun.sourceOp, storedRun.resultIndex, storedRun.targetClass};
-    packedRunsByProducerResultClass[lookupKey].push_back(runIndex);
-  }
-
-  void recordIndexedBatchRun(IndexedBatchRunValue run) { indexedBatchRuns.push_back(std::move(run)); }
-  void recordDeferredExact(ProducerKey key, ClassId classId, llvm::ArrayRef<CommunicationExchangeId> exchangeIds) {
-    llvm::SmallVector<CommunicationExchangeId, 4>& bucket = exactExchangeIds[key][classId];
-    for (CommunicationExchangeId id : exchangeIds)
-      if (!llvm::is_contained(bucket, id))
-        bucket.push_back(id);
-  }
-  llvm::ArrayRef<CommunicationExchangeId> getDeferredExactExchangeIds(ProducerKey key, ClassId classId) const {
-    auto producerIt = exactExchangeIds.find(key);
-    if (producerIt == exactExchangeIds.end())
-      return {};
-    auto classIt = producerIt->second.find(classId);
-    if (classIt == producerIt->second.end())
-      return {};
-    return classIt->second;
-  }
-  bool hasDeferredExact(ProducerKey key, ClassId classId) const {
-    return !getDeferredExactExchangeIds(key, classId).empty();
-  }
-
-  std::optional<mlir::Value> lookupExact(ProducerKey key, ClassId classId) const;
-  std::optional<mlir::Value> lookup(MaterializerState& state, ProducerKey key, ClassId classId);
-  IndexedBatchRunValue* lookupIndexedBatchRun(ProducerKey key, ClassId classId);
-
-  llvm::ArrayRef<size_t> getPackedRunIndicesForWholeBatch(WholeBatchAssemblyLookupKey key) const {
-    auto it = packedRunsByProducerResultClass.find(key);
-    if (it == packedRunsByProducerResultClass.end())
-      return {};
-    return it->second;
-  }
-
-  llvm::ArrayRef<ExactBatchFragmentRecord> getExactFragmentsForWholeBatch(WholeBatchAssemblyLookupKey key) const {
-    auto it = exactBatchFragmentsByProducerResultClass.find(key);
-    if (it == exactBatchFragmentsByProducerResultClass.end())
-      return {};
-    return it->second;
-  }
-
-  PackedScalarRunValue& getPackedRun(size_t index) { return packedScalarRuns[index]; }
-
-private:
-  std::optional<mlir::Value> lookupPackedRun(MaterializerState& state, ProducerKey key, ClassId classId);
-
-  llvm::DenseMap<ProducerKey, llvm::DenseMap<ClassId, mlir::Value>, ProducerKeyInfo> exactValues;
-  llvm::DenseMap<ProducerKey,
-                 llvm::DenseMap<ClassId, llvm::SmallVector<CommunicationExchangeId, 4>>,
-                 ProducerKeyInfo>
-    exactExchangeIds;
-  llvm::SmallVector<PackedScalarRunValue, 8> packedScalarRuns;
-  llvm::SmallVector<IndexedBatchRunValue, 8> indexedBatchRuns;
-  llvm::DenseMap<WholeBatchAssemblyLookupKey,
-                 llvm::SmallVector<ExactBatchFragmentRecord, 16>,
-                 WholeBatchAssemblyLookupKeyInfo>
-    exactBatchFragmentsByProducerResultClass;
-  llvm::DenseMap<WholeBatchAssemblyLookupKey, llvm::SmallVector<size_t, 16>, WholeBatchAssemblyLookupKeyInfo>
-    packedRunsByProducerResultClass;
-};
-
-struct MaterializerState {
-  mlir::func::FuncOp func;
-  const MergeScheduleResult& schedule;
-  mlir::IRRewriter rewriter;
-  mlir::OperationFolder constantFolder;
-  int64_t& nextChannelId;
-  llvm::SmallVector<MaterializedClass, 8> classes;
-  llvm::DenseMap<CpuId, ClassId> cpuToClass;
-  llvm::DenseMap<CpuId, llvm::SmallVector<ComputeInstance, 32>> logicalInstancesByCpu;
-  llvm::DenseMap<ComputeInstance, LogicalSlotRange> scheduledInstanceToLogicalSlots;
-  llvm::DenseMap<ComputeInstance, ComputeInstance> logicalInstanceToScheduledChunk;
-  llvm::DenseSet<ClassSlotKey> materializedLogicalSlots;
-
-  llvm::DenseMap<ProducerKey, llvm::SmallVector<ClassId, 4>, ProducerKeyInfo> producerDestClasses;
-  llvm::DenseMap<ProducerKey, llvm::SmallVector<ClassId, 4>, ProducerKeyInfo> ordinaryProducerDestClasses;
-  llvm::DenseMap<SameClassConsumerLookupKey, llvm::SmallVector<ProducerKey, 4>, SameClassConsumerLookupKeyInfo>
-    sameClassConsumerIndex;
-  llvm::DenseMap<ProjectedBatchInputKey, AffineProjectedInputSliceMatch, ProjectedBatchInputKeyInfo>
-    projectedInputMatches;
-  llvm::DenseSet<ProjectedBatchInputKey, ProjectedBatchInputKeyInfo> nonProjectedInputs;
-  llvm::DenseMap<mlir::Value, bool> liveExternalUseCache;
-  llvm::DenseMap<mlir::Operation*, llvm::SmallVector<mlir::Type, 4>> batchOutputFragmentTypesCache;
-  llvm::DenseMap<ComputeInstance, llvm::SmallVector<mlir::Value, 4>, llvm::DenseMapInfo<ComputeInstance>>
-    computeInstanceOutputsCache;
-  llvm::DenseMap<ProducerKey, llvm::DenseMap<ClassId, ProjectedTransferDescriptor>, ProducerKeyInfo>
-    projectedTransfers;
-  llvm::DenseMap<mlir::Operation*, llvm::DenseMap<ClassId, ProjectedExtractReplacement>>
-    projectedExtractReplacements;
-  llvm::DenseMap<mlir::Operation*, llvm::DenseMap<ClassId, ProjectedInputTransferPlan>>
-    projectedInputTransferPlans;
-  CommunicationPlan communicationPlan;
-  llvm::SmallVector<CommunicationExchangeEmissionState, 16> communicationExchangeStates;
-  llvm::SmallVector<ProducerKey, 8> deferredProducerValues;
-  AvailableValueStore availableValues;
-  llvm::DenseMap<mlir::Value, mlir::Value> hostReplacements;
-  llvm::DenseMap<mlir::Value, ClassId> hostOutputOwners;
-  llvm::SmallVector<PendingProjectedHostOutputFragment, 32> pendingProjectedHostOutputFragments;
-  llvm::DenseSet<mlir::Operation*> oldComputeOps;
-
-  MaterializerState(mlir::func::FuncOp func, const MergeScheduleResult& schedule, int64_t& nextChannelId)
-  : func(func),
-    schedule(schedule),
-    rewriter(func.getContext()),
-    constantFolder(func.getContext()),
-    nextChannelId(nextChannelId) {}
-};
-
-} // namespace onnx_mlir::spatial

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

@@ -1,398 +0,0 @@
-#include "mlir/Analysis/TopologicalSortUtils.h"
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/IR/IRMapping.h"
-#include "mlir/IR/Location.h"
-#include "mlir/IR/PatternMatch.h"
-#include "mlir/IR/Region.h"
-#include "mlir/IR/Value.h"
-#include "mlir/IR/ValueRange.h"
-#include "mlir/Pass/Pass.h"
-#include "mlir/Support/LLVM.h"
-
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/SmallSet.h"
-#include "llvm/ADT/SmallVector.h"
-#include "llvm/Support/raw_os_ostream.h"
-
-#include <algorithm>
-#include <cstddef>
-#include <cstdint>
-#include <fstream>
-#include <memory>
-#include <optional>
-#include <utility>
-#include <vector>
-
-#include "MaterializeMergeSchedule.hpp"
-#include "SpatialDataflowCsvExporter.hpp"
-#include "Scheduling/ComputeGraph.hpp"
-#include "Scheduling/ComputeInstanceUtils.hpp"
-#include "Scheduling/MergeSchedulingAnalysis.hpp"
-#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
-#include "src/Accelerators/PIM/Common/IR/CompactAsmUtils.hpp"
-#include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
-#include "src/Accelerators/PIM/Common/Support/ReportUtils.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-
-using namespace mlir;
-
-namespace onnx_mlir {
-namespace {
-using namespace onnx_mlir::compact_asm;
-using SpatCompute = spatial::SpatGraphCompute;
-using SpatComputeBatch = spatial::SpatGraphComputeBatch;
-
-bool isTrivialSerialMergeCandidate(SpatCompute compute) {
-  if (!compute->hasOneUse())
-    return false;
-
-  auto& use = *compute->getUses().begin();
-  auto user = dyn_cast<SpatCompute>(use.getOwner());
-  return user && user.getInputs().size() == 1 && use.getOperandNumber() >= user.getWeights().size();
-}
-
-SmallVector<size_t> appendMissingWeightsAndBuildIndexMap(SmallVectorImpl<Value>& targetWeights,
-                                                         ValueRange sourceWeights) {
-  DenseMap<Value, SmallVector<size_t, 4>> targetWeightIndices;
-  for (auto [weightIndex, weight] : llvm::enumerate(targetWeights))
-    targetWeightIndices[weight].push_back(weightIndex);
-
-  DenseMap<Value, size_t> usedSourceWeightOccurrences;
-  SmallVector<size_t> sourceToTargetIndex;
-  sourceToTargetIndex.reserve(sourceWeights.size());
-  for (Value weight : sourceWeights) {
-    size_t occurrence = usedSourceWeightOccurrences[weight]++;
-    auto& matchingIndices = targetWeightIndices[weight];
-    if (occurrence >= matchingIndices.size()) {
-      size_t newIndex = targetWeights.size();
-      targetWeights.push_back(weight);
-      matchingIndices.push_back(newIndex);
-      sourceToTargetIndex.push_back(newIndex);
-      continue;
-    }
-    sourceToTargetIndex.push_back(matchingIndices[occurrence]);
-  }
-  return sourceToTargetIndex;
-}
-
-void mergeTriviallyConnectedComputes(func::FuncOp funcOp) {
-  Location loc = funcOp.getLoc();
-  IRRewriter rewriter(funcOp->getContext());
-  SmallVector<SpatCompute> trivialComputes;
-  llvm::SmallSet<SpatCompute, 8> toErase;
-
-  for (auto compute : funcOp.getOps<SpatCompute>())
-    if (isTrivialSerialMergeCandidate(compute))
-      trivialComputes.push_back(compute);
-
-  while (!trivialComputes.empty()) {
-    auto compute = trivialComputes.front();
-
-    if (compute.use_empty()) {
-      std::swap(trivialComputes.front(), trivialComputes.back());
-      trivialComputes.pop_back();
-      continue;
-    }
-
-    auto& computeUse = *compute->getUses().begin();
-    auto child = cast<SpatCompute>(computeUse.getOwner());
-    auto usedResult = cast<OpResult>(computeUse.get()).getResultNumber();
-    auto childInputIndex = computeUse.getOperandNumber() - child.getWeights().size();
-
-    rewriter.setInsertionPointAfter(compute.getOperation());
-    SmallVector<Value> mergedWeights(compute.getWeights().begin(), compute.getWeights().end());
-    SmallVector<size_t> childWeightToNewIndex = appendMissingWeightsAndBuildIndexMap(mergedWeights, child.getWeights());
-    SmallVector<Value> mergedInputs(compute.getInputs().begin(), compute.getInputs().end());
-    auto newCompute = SpatCompute::create(rewriter, loc, child.getResultTypes(), mergedWeights, mergedInputs);
-    Block* newBody = rewriter.createBlock(&newCompute.getBodyRegion());
-    for (Value weight : mergedWeights)
-      newBody->addArgument(weight.getType(), loc);
-    for (Value input : mergedInputs)
-      newBody->addArgument(input.getType(), loc);
-
-    IRMapping mapper;
-    for (auto [weightIndex, _] : llvm::enumerate(compute.getWeights())) {
-      auto oldWeightArg = compute.getWeightArgument(weightIndex);
-      auto newWeightArg = newCompute.getWeightArgument(weightIndex);
-      assert(oldWeightArg && newWeightArg && "expected compute weight block arguments");
-      mapper.map(*oldWeightArg, *newWeightArg);
-    }
-    for (auto [inputIndex, _] : llvm::enumerate(compute.getInputs())) {
-      auto oldInputArg = compute.getInputArgument(inputIndex);
-      auto newInputArg = newCompute.getInputArgument(inputIndex);
-      assert(oldInputArg && newInputArg && "expected compute input block arguments");
-      mapper.map(*oldInputArg, *newInputArg);
-    }
-    for (auto [oldIndex, weight] : llvm::enumerate(child.getWeights())) {
-      auto oldWeightArg = child.getWeightArgument(oldIndex);
-      auto newWeightArg = newCompute.getWeightArgument(childWeightToNewIndex[oldIndex]);
-      assert(oldWeightArg && newWeightArg && "expected child compute weight block arguments");
-      mapper.map(*oldWeightArg, *newWeightArg);
-    }
-
-    rewriter.setInsertionPointToEnd(newBody);
-    auto computeYield = cast<spatial::SpatYieldOp>(compute.getBody().front().getTerminator());
-    for (Operation& op : compute.getBody().front().without_terminator())
-      rewriter.clone(op, mapper);
-    auto childInputArg = child.getInputArgument(childInputIndex);
-    assert(childInputArg && "expected child compute input block argument");
-    mapper.map(*childInputArg, mapper.lookupOrDefault(computeYield.getOperand(usedResult)));
-
-    rewriter.setInsertionPointToEnd(newBody);
-    for (auto& op : child.getBody().front())
-      rewriter.clone(op, mapper);
-
-    child.replaceAllUsesWith(newCompute);
-    toErase.insert(child);
-
-    std::swap(trivialComputes.front(), trivialComputes.back());
-    trivialComputes.pop_back();
-    toErase.insert(compute);
-
-    if (isTrivialSerialMergeCandidate(newCompute))
-      trivialComputes.push_back(newCompute);
-  }
-
-  for (auto compute : toErase) {
-    for (Value result : compute->getResults())
-      result.dropAllUses();
-    compute.erase();
-  }
-}
-
-void generateReport(func::FuncOp funcOp, const std::string& name, size_t usedCpuCount = 0) {
-  std::fstream file = openReportFile(name);
-  if (!file.is_open())
-    return;
-  llvm::raw_os_ostream os(file);
-
-  struct ReportRow {
-    uint64_t id = 0;
-    uint64_t logicalComputeCount = 0;
-    uint64_t crossbarCount = 0;
-    uint64_t instructionCount = 0;
-    bool isRebatched = false;
-    SmallVector<int32_t> coreIds;
-  };
-
-  //TODO Used for report refactor
-  struct CollectorConcatRow {
-    uint64_t computeId = 0;
-    int32_t coreId = -1;
-    uint64_t operandCount = 0;
-  };
-
-  uint64_t totalComputeOps = 0;
-  uint64_t totalLogicalComputes = 0;
-  uint64_t totalBatchComputeOps = 0;
-  uint64_t totalInstructionCount = 0;
-  uint64_t totalCrossbarCount = 0;
-  uint64_t nextBatchId = 0;
-  //TODO Used for report refactor
-  std::vector<ReportRow> collectedData;
-  //TODO Used for report refactor
-  std::vector<CollectorConcatRow> collectorConcatRows;
-
-  auto getPerInstanceCrossbarCount = [&](Operation* op) -> uint64_t {
-    return static_cast<uint64_t>(spatial::collectDistinctCrossbarWeights(op).size());
-  };
-
-  for (Operation& op : funcOp.getBody().front()) {
-    if (auto spatCompute = dyn_cast<spatial::SpatScheduledCompute>(&op)) {
-      uint64_t numInst = spatial::countComputeBodyInstructions(spatCompute.getBody());
-      uint64_t perInstanceCrossbarCount = getPerInstanceCrossbarCount(spatCompute.getOperation());
-      SmallVector<int32_t> coreIds;
-      auto coreId = getOptionalScheduledCoreId(spatCompute, "merge compute core id");
-      if (failed(coreId))
-        return;
-      if (*coreId)
-        coreIds.push_back(**coreId);
-      uint64_t computeId = totalComputeOps++;
-      collectedData.push_back({computeId, 1, perInstanceCrossbarCount, numInst, false, coreIds});
-      uint64_t maxConcatOperands = 0;
-      spatCompute.getBody().walk([&](spatial::SpatConcatOp concatOp) {
-        maxConcatOperands = std::max<uint64_t>(maxConcatOperands, concatOp.getInputs().size());
-      });
-      //TODO 128 is a magic number
-      if (maxConcatOperands >= 128 && !coreIds.empty())
-        collectorConcatRows.push_back({computeId, coreIds.front(), maxConcatOperands});
-      totalLogicalComputes += 1;
-      totalInstructionCount += numInst;
-      totalCrossbarCount += perInstanceCrossbarCount;
-      continue;
-    }
-    if (auto batch = dyn_cast<spatial::SpatScheduledComputeBatch>(&op)) {
-      uint64_t numInst = spatial::countComputeBodyInstructions(batch.getBody());
-      uint64_t logicalCount = static_cast<uint64_t>(batch.getLaneCount());
-      uint64_t perInstanceCrossbarCount = getPerInstanceCrossbarCount(batch.getOperation());
-      SmallVector<int32_t> coreIds;
-      auto optionalCoreIds = getOptionalScheduledBatchCoreIds(batch, "merge compute_batch core id");
-      if (failed(optionalCoreIds))
-        return;
-      if (*optionalCoreIds)
-        coreIds = std::move(**optionalCoreIds);
-      collectedData.push_back(
-        {nextBatchId++, logicalCount, perInstanceCrossbarCount * logicalCount, numInst, true, coreIds});
-      totalComputeOps += 1;
-      totalLogicalComputes += logicalCount;
-      totalBatchComputeOps += 1;
-      totalInstructionCount += numInst * logicalCount;
-      totalCrossbarCount += perInstanceCrossbarCount * logicalCount;
-    }
-  }
-
-  llvm::SmallVector<ReportField, 6> totalFields = {
-    {"Used cores",                            std::to_string(usedCpuCount)         },
-    {"Number of top-level compute ops",       std::to_string(totalComputeOps)      },
-    {"Number of logical computes",            std::to_string(totalLogicalComputes) },
-    {"Number of top-level batch compute ops", std::to_string(totalBatchComputeOps) },
-    {"Number of instructions",                std::to_string(totalInstructionCount)},
-    {"Number of used crossbars",              std::to_string(totalCrossbarCount)   }
-  };
-  printReportTotalsBlock(os, totalFields);
-  if (!collectedData.empty() || !collectorConcatRows.empty())
-    os << "\n";
-
-  if (!collectorConcatRows.empty()) {
-    os << "Collector concat materialization:\n";
-    for (const CollectorConcatRow& row : collectorConcatRows)
-      os << "\tmaterialization_kind = single_collector_concat, compute = " << row.computeId
-         << ", concat_operand_count = " << row.operandCount << ", collector_core = " << row.coreId << "\n";
-    os << "\n";
-  }
-
-  sortReportEntriesByFirstCore(collectedData);
-
-  for (uint64_t cI = 0; cI < totalComputeOps; ++cI) {
-    uint64_t lastIndex = cI;
-    ReportRow current = collectedData[cI];
-
-    for (uint64_t nI = cI + 1; nI < totalComputeOps; ++nI) {
-      ReportRow next = collectedData[nI];
-      if (current.isRebatched == next.isRebatched && current.crossbarCount == next.crossbarCount
-          && current.instructionCount == next.instructionCount
-          && current.logicalComputeCount == next.logicalComputeCount)
-        lastIndex = nI;
-      else
-        break;
-    }
-
-    if (current.isRebatched) {
-      os << "Batch ";
-      for (uint64_t index = cI; index <= lastIndex; ++index) {
-        if (index != cI)
-          os << ",\n      ";
-        os << collectedData[index].id << " (cores ";
-        if (collectedData[index].coreIds.empty())
-          os << "unknown";
-        else
-          printCompressedIntegerEntries(os, ArrayRef<int32_t>(collectedData[index].coreIds));
-        os << ")";
-      }
-    }
-    else {
-      os << "Compute ";
-      SmallVector<uint64_t> opIds;
-      opIds.reserve(lastIndex - cI + 1);
-      for (uint64_t index = cI; index <= lastIndex; ++index)
-        opIds.push_back(collectedData[index].id);
-      printCompressedIntegerEntries(os, ArrayRef<uint64_t>(opIds));
-    }
-
-    os << ":\n";
-    uint64_t perCoreLogicalComputeCount = current.isRebatched ? 1 : current.logicalComputeCount;
-    uint64_t perCoreInstructionCount = current.instructionCount;
-    uint64_t perCoreCrossbarCount =
-      current.logicalComputeCount == 0 ? 0 : current.crossbarCount / current.logicalComputeCount;
-    uint64_t totalEntryInstructionCount = current.instructionCount * current.logicalComputeCount;
-
-    llvm::SmallVector<ReportField, 3> perCoreFields = {
-      {"Number of logical computes", std::to_string(perCoreLogicalComputeCount)},
-      {"Number of instructions",     std::to_string(perCoreInstructionCount)   },
-      {"Number of used crossbars",   std::to_string(perCoreCrossbarCount)      }
-    };
-    if (current.isRebatched) {
-      llvm::SmallVector<ReportField, 3> totalEntryFields = {
-        {"Number of logical computes", std::to_string(current.logicalComputeCount)},
-        {"Number of instructions",     std::to_string(totalEntryInstructionCount) },
-        {"Number of used crossbars",   std::to_string(current.crossbarCount)      }
-      };
-      printReportPerCoreAndTotalFields(os, perCoreFields, totalEntryFields);
-    }
-    else {
-      printReportFlatFields(os, perCoreFields);
-    }
-    printReportEntrySeparator(os, lastIndex + 1 < totalComputeOps);
-    cI = lastIndex;
-  }
-
-  os.flush();
-  file.close();
-}
-
-struct MergeComputeNodesPass : PassWrapper<MergeComputeNodesPass, OperationPass<func::FuncOp>> {
-
-private:
-  int64_t nextChannelId = 0;
-
-public:
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MergeComputeNodesPass)
-
-  StringRef getArgument() const override { return "pim-merge-compute-nodes-pass"; }
-  StringRef getDescription() const override {
-    return "Merge Spatial-Compute-Nodes in order to reduce the total "
-           "execution time";
-  }
-
-  LogicalResult initialize(MLIRContext* context) override { return success(); }
-
-  void runOnOperation() override {
-    func::FuncOp func = getOperation();
-    if (failed(verifyLogicalSpatialGraphInvariants(func))) {
-      func.emitOpError("logical Spatial graph verification failed at the start of MergeComputeNodes");
-      signalPassFailure();
-      return;
-    }
-    mergeTriviallyConnectedComputes(func);
-    if (failed(verifyLogicalSpatialGraphInvariants(func))) {
-      func.emitOpError("logical Spatial graph verification failed after trivial merge simplification");
-      signalPassFailure();
-      return;
-    }
-
-    const spatial::MergeScheduleResult* analysisResult = nullptr;
-    analysisResult = &getAnalysis<spatial::MergeSchedulingAnalysis>().getResult();
-    spatial::SpatialDataflowExportStage exportMode = spatial::getSpatialDataflowExportStage();
-    if (failed(spatial::MergeScheduleMaterializer().run(func, *analysisResult, nextChannelId))) {
-      signalPassFailure();
-      return;
-    }
-
-    if (!sortTopologically(&func.getBody().front())) {
-      func.emitOpError("failed to topologically order merged Spatial IR");
-      signalPassFailure();
-      return;
-    }
-    if (failed(verifyScheduledSpatialInvariants(func))) {
-      func.emitOpError("scheduled Spatial verification failed after merge materialization");
-      signalPassFailure();
-      return;
-    }
-    if (spatial::shouldExportSpatialDataflowStage(exportMode, spatial::SpatialDataflowExportStage::Post)
-        && failed(spatial::exportSpatialDataflowCsvPost(func))) {
-      signalPassFailure();
-      return;
-    }
-    dumpModule(cast<ModuleOp>(func->getParentOp()), "spatial2_merged");
-    generateReport(func, "spatial_merge_report", analysisResult->cpuToLastComputeMap.size());
-  }
-};
-
-} // namespace
-
-std::unique_ptr<Pass> createMergeComputeNodesPass() { return std::make_unique<MergeComputeNodesPass>(); }
-
-} // namespace onnx_mlir

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeMessages.hpp

@@ -1,67 +0,0 @@
-#pragma once
-
-#include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/SmallVector.h"
-#include "llvm/ADT/StringRef.h"
-
-#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
-
-namespace onnx_mlir::spatial {
-
-using CpuId = size_t;
-
-inline mlir::FailureOr<int32_t> getCheckedCoreId(mlir::Operation* anchor, CpuId cpu, llvm::StringRef fieldName) {
-  return pim::checkedI32(static_cast<uint64_t>(cpu), anchor, fieldName);
-}
-
-inline mlir::FailureOr<llvm::SmallVector<int32_t, 8>>
-getCheckedCoreIds(mlir::Operation* anchor, llvm::ArrayRef<CpuId> cpus, llvm::StringRef fieldName) {
-  llvm::SmallVector<int32_t, 8> coreIds;
-  coreIds.reserve(cpus.size());
-  for (CpuId cpu : cpus) {
-    auto checkedCoreId = getCheckedCoreId(anchor, cpu, fieldName);
-    if (mlir::failed(checkedCoreId))
-      return mlir::failure();
-    coreIds.push_back(*checkedCoreId);
-  }
-  return coreIds;
-}
-
-struct MessageVector {
-  llvm::SmallVector<int64_t, 16> channelIds;
-  llvm::SmallVector<int32_t, 16> sourceCoreIds;
-  llvm::SmallVector<int32_t, 16> targetCoreIds;
-
-  size_t size() const { return channelIds.size(); }
-  bool empty() const { return channelIds.empty(); }
-
-  mlir::LogicalResult verify(mlir::Operation* anchor) const {
-    if (channelIds.size() != sourceCoreIds.size() || channelIds.size() != targetCoreIds.size())
-      return anchor->emitError("message metadata is inconsistent");
-    return mlir::success();
-  }
-
-  void append(int64_t channelId, int32_t sourceCoreId, int32_t targetCoreId) {
-    channelIds.push_back(channelId);
-    sourceCoreIds.push_back(sourceCoreId);
-    targetCoreIds.push_back(targetCoreId);
-  }
-
-  void append(llvm::ArrayRef<int64_t> channels, llvm::ArrayRef<int32_t> sources, llvm::ArrayRef<int32_t> targets) {
-    assert(channels.size() == sources.size() && "channel/source count mismatch");
-    assert(channels.size() == targets.size() && "channel/target count mismatch");
-    llvm::append_range(channelIds, channels);
-    llvm::append_range(sourceCoreIds, sources);
-    llvm::append_range(targetCoreIds, targets);
-  }
-
-  MessageVector slice(size_t offset, size_t count) const {
-    MessageVector result;
-    result.append(llvm::ArrayRef<int64_t>(channelIds).slice(offset, count),
-                  llvm::ArrayRef<int32_t>(sourceCoreIds).slice(offset, count),
-                  llvm::ArrayRef<int32_t>(targetCoreIds).slice(offset, count));
-    return result;
-  }
-};
-
-} // namespace onnx_mlir::spatial

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeScheduleKeys.hpp

@@ -1,134 +0,0 @@
-#pragma once
-
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/SmallVector.h"
-
-#include <cstddef>
-#include <cstdint>
-#include <limits>
-#include <utility>
-
-#include "Scheduling/ComputeInstanceUtils.hpp"
-
-namespace onnx_mlir::spatial {
-
-using ClassId = size_t;
-using SlotId = size_t;
-
-struct ProducerKey {
-  ComputeInstance instance;
-  size_t resultIndex = 0;
-
-  bool operator==(const ProducerKey& other) const {
-    return instance == other.instance && resultIndex == other.resultIndex;
-  }
-};
-
-struct ProducerKeyInfo {
-  static ProducerKey getEmptyKey() {
-    return {llvm::DenseMapInfo<ComputeInstance>::getEmptyKey(), std::numeric_limits<size_t>::max()};
-  }
-
-  static ProducerKey getTombstoneKey() {
-    return {llvm::DenseMapInfo<ComputeInstance>::getTombstoneKey(), std::numeric_limits<size_t>::max()};
-  }
-
-  static unsigned getHashValue(const ProducerKey& key) {
-    return llvm::hash_combine(llvm::DenseMapInfo<ComputeInstance>::getHashValue(key.instance), key.resultIndex);
-  }
-
-  static bool isEqual(const ProducerKey& lhs, const ProducerKey& rhs) { return lhs == rhs; }
-};
-
-struct SameClassConsumerLookupKey {
-  mlir::Operation* sourceOp = nullptr;
-  size_t resultIndex = 0;
-  ClassId classId = 0;
-
-  bool operator==(const SameClassConsumerLookupKey& other) const {
-    return sourceOp == other.sourceOp && resultIndex == other.resultIndex && classId == other.classId;
-  }
-};
-
-struct SameClassConsumerLookupKeyInfo {
-  static SameClassConsumerLookupKey getEmptyKey() {
-    return {llvm::DenseMapInfo<mlir::Operation*>::getEmptyKey(), std::numeric_limits<size_t>::max(),
-            std::numeric_limits<ClassId>::max()};
-  }
-
-  static SameClassConsumerLookupKey getTombstoneKey() {
-    return {llvm::DenseMapInfo<mlir::Operation*>::getTombstoneKey(), std::numeric_limits<size_t>::max(),
-            std::numeric_limits<ClassId>::max()};
-  }
-
-  static unsigned getHashValue(const SameClassConsumerLookupKey& key) {
-    return llvm::hash_combine(llvm::DenseMapInfo<mlir::Operation*>::getHashValue(key.sourceOp),
-                              key.resultIndex,
-                              key.classId);
-  }
-
-  static bool isEqual(const SameClassConsumerLookupKey& lhs, const SameClassConsumerLookupKey& rhs) {
-    return lhs == rhs;
-  }
-};
-
-struct WholeBatchAssemblyLookupKey {
-  mlir::Operation* sourceOp = nullptr;
-  size_t resultIndex = 0;
-  ClassId classId = 0;
-
-  bool operator==(const WholeBatchAssemblyLookupKey& other) const {
-    return sourceOp == other.sourceOp && resultIndex == other.resultIndex && classId == other.classId;
-  }
-};
-
-struct WholeBatchAssemblyLookupKeyInfo {
-  static WholeBatchAssemblyLookupKey getEmptyKey() {
-    return {llvm::DenseMapInfo<mlir::Operation*>::getEmptyKey(), std::numeric_limits<size_t>::max(),
-            std::numeric_limits<ClassId>::max()};
-  }
-
-  static WholeBatchAssemblyLookupKey getTombstoneKey() {
-    return {llvm::DenseMapInfo<mlir::Operation*>::getTombstoneKey(), std::numeric_limits<size_t>::max(),
-            std::numeric_limits<ClassId>::max()};
-  }
-
-  static unsigned getHashValue(const WholeBatchAssemblyLookupKey& key) {
-    return llvm::hash_combine(llvm::DenseMapInfo<mlir::Operation*>::getHashValue(key.sourceOp),
-                              key.resultIndex,
-                              key.classId);
-  }
-
-  static bool isEqual(const WholeBatchAssemblyLookupKey& lhs, const WholeBatchAssemblyLookupKey& rhs) {
-    return lhs == rhs;
-  }
-};
-
-using ClassSlotKey = std::pair<ClassId, SlotId>;
-
-struct ProjectedBatchInputKey {
-  mlir::Operation* consumerOp = nullptr;
-  unsigned inputIndex = 0;
-
-  bool operator==(const ProjectedBatchInputKey& other) const {
-    return consumerOp == other.consumerOp && inputIndex == other.inputIndex;
-  }
-};
-
-struct ProjectedBatchInputKeyInfo {
-  static ProjectedBatchInputKey getEmptyKey() {
-    return {llvm::DenseMapInfo<mlir::Operation*>::getEmptyKey(), std::numeric_limits<unsigned>::max()};
-  }
-
-  static ProjectedBatchInputKey getTombstoneKey() {
-    return {llvm::DenseMapInfo<mlir::Operation*>::getTombstoneKey(), std::numeric_limits<unsigned>::max()};
-  }
-
-  static unsigned getHashValue(const ProjectedBatchInputKey& key) {
-    return llvm::hash_combine(key.consumerOp, key.inputIndex);
-  }
-
-  static bool isEqual(const ProjectedBatchInputKey& lhs, const ProjectedBatchInputKey& rhs) { return lhs == rhs; }
-};
-
-} // namespace onnx_mlir::spatial

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

@@ -1,104 +0,0 @@
-#include "ProjectedFragments.hpp"
-
-#include "mlir/IR/BuiltinTypes.h"
-
-namespace onnx_mlir::spatial {
-
-static mlir::FailureOr<mlir::RankedTensorType> getPackedBatchTensorType(mlir::Type laneType, size_t laneCount) {
-  auto tensorType = mlir::dyn_cast<mlir::RankedTensorType>(laneType);
-  if (!tensorType || !tensorType.hasStaticShape() || tensorType.getRank() == 0)
-    return mlir::failure();
-
-  llvm::SmallVector<int64_t, 4> shape(tensorType.getShape());
-  shape[0] *= static_cast<int64_t>(laneCount);
-  return mlir::RankedTensorType::get(shape, tensorType.getElementType(), tensorType.getEncoding());
-}
-
-unsigned getProjectedFragmentsPerLogicalSlot(llvm::ArrayRef<int64_t> loopTripCounts) {
-  unsigned fragmentsPerLogicalSlot = 1;
-  for (int64_t tripCount : loopTripCounts) {
-    assert(tripCount > 0 && "projected loop trip counts must be positive");
-    fragmentsPerLogicalSlot *= static_cast<unsigned>(tripCount);
-  }
-  return fragmentsPerLogicalSlot;
-}
-
-mlir::LogicalResult verifyProjectedFragmentLayout(mlir::Operation* anchor, const ProjectedFragmentLayout& layout) {
-  if (!layout.fragmentType || layout.fragmentShape.empty())
-    return anchor->emitError("projected fragment layout is missing fragment type metadata");
-  if (layout.fragmentShape.size() != static_cast<size_t>(layout.fragmentType.getRank()))
-    return anchor->emitError("projected fragment layout rank does not match fragment type");
-  if (layout.payloadFragmentCount == 0 || layout.fragmentsPerLogicalSlot == 0)
-    return anchor->emitError("projected fragment layout has an invalid fragment count");
-  if (layout.payloadFragmentCount % layout.fragmentsPerLogicalSlot != 0)
-    return anchor->emitError("projected fragment layout payload fragment count is incompatible with logical slots");
-  return mlir::success();
-}
-
-mlir::FailureOr<mlir::RankedTensorType>
-getProjectedPayloadType(mlir::Operation* anchor, mlir::RankedTensorType fragmentType, unsigned payloadFragmentCount) {
-  auto packedType = getPackedBatchTensorType(fragmentType, payloadFragmentCount);
-  if (mlir::failed(packedType)) {
-    anchor->emitError("cannot create projected payload type");
-    return mlir::failure();
-  }
-  return *packedType;
-}
-
-llvm::SmallVector<llvm::SmallVector<int64_t, 16>, 4>
-buildProjectedFragmentOffsetsByDim(llvm::ArrayRef<llvm::SmallVector<int64_t, 4>> fragmentOffsets, size_t rank) {
-  llvm::SmallVector<llvm::SmallVector<int64_t, 16>, 4> fragmentOffsetsByDim(rank);
-  for (llvm::ArrayRef<int64_t> offsets : fragmentOffsets) {
-    assert(offsets.size() == rank && "projected offset rank mismatch");
-    for (size_t dim = 0; dim < rank; ++dim)
-      fragmentOffsetsByDim[dim].push_back(offsets[dim]);
-  }
-  return fragmentOffsetsByDim;
-}
-
-mlir::LogicalResult verifyProjectedTransferDescriptor(mlir::Operation* anchor,
-                                                      const ProjectedTransferDescriptor& descriptor) {
-  if (mlir::failed(verifyProjectedFragmentLayout(anchor, descriptor.layout)))
-    return mlir::failure();
-  if (!descriptor.payloadType)
-    return anchor->emitError("projected transfer descriptor is missing payload type");
-  if (descriptor.fragmentOffsets.empty())
-    return anchor->emitError("projected transfer descriptor expected at least one fragment offset");
-  if (descriptor.fragmentOffsetsByDim.size() != descriptor.layout.fragmentShape.size())
-    return anchor->emitError("projected transfer descriptor dimension-major offsets are inconsistent");
-  for (llvm::ArrayRef<int64_t> dimOffsets : descriptor.fragmentOffsetsByDim)
-    if (dimOffsets.size() != descriptor.fragmentOffsets.size())
-      return anchor->emitError("projected transfer descriptor dimension-major offsets are inconsistent");
-  for (llvm::ArrayRef<int64_t> offsets : descriptor.fragmentOffsets)
-    if (offsets.size() != descriptor.layout.fragmentShape.size())
-      return anchor->emitError("projected transfer offset rank does not match fragment rank");
-  return mlir::success();
-}
-
-mlir::LogicalResult verifyProjectedSendDescriptor(mlir::Operation* anchor,
-                                                  const ProjectedTransferDescriptor& descriptor,
-                                                  const MessageVector& messages) {
-  if (mlir::failed(verifyProjectedTransferDescriptor(anchor, descriptor)))
-    return mlir::failure();
-  if (messages.size() * descriptor.layout.payloadFragmentCount != descriptor.fragmentOffsets.size())
-    return anchor->emitError("projected send descriptor metadata is inconsistent");
-  return mlir::success();
-}
-
-mlir::LogicalResult finalizeProjectedTransferDescriptor(mlir::Operation* anchor,
-                                                        ProjectedTransferDescriptor& descriptor) {
-  descriptor.fragmentOffsetsByDim =
-    buildProjectedFragmentOffsetsByDim(descriptor.fragmentOffsets, descriptor.layout.fragmentShape.size());
-
-  auto payloadType =
-    getProjectedPayloadType(anchor, descriptor.layout.fragmentType, descriptor.layout.payloadFragmentCount);
-  if (mlir::failed(payloadType))
-    return mlir::failure();
-  if (descriptor.payloadType && descriptor.payloadType != *payloadType)
-    return anchor->emitError("projected transfer descriptor payload type does not match projected layout");
-  descriptor.payloadType = *payloadType;
-
-  return verifyProjectedTransferDescriptor(anchor, descriptor);
-}
-
-} // namespace onnx_mlir::spatial

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ProjectedFragments.hpp

@@ -1,105 +0,0 @@
-#pragma once
-
-#include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/IR/BuiltinAttributes.h"
-#include "mlir/IR/BuiltinTypes.h"
-#include "mlir/IR/Value.h"
-#include "mlir/IR/ValueRange.h"
-
-#include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/SmallVector.h"
-
-#include <cstdint>
-
-#include "MergeMessages.hpp"
-#include "MergeScheduleKeys.hpp"
-
-namespace onnx_mlir::spatial {
-
-struct ProjectedFragmentLayout {
-  mlir::RankedTensorType fragmentType;
-  llvm::SmallVector<int64_t, 4> fragmentShape;
-  unsigned fragmentsPerLogicalSlot = 1;
-  unsigned payloadFragmentCount = 1;
-  llvm::SmallVector<int64_t, 4> loopLowerBounds;
-  llvm::SmallVector<int64_t, 4> loopSteps;
-  llvm::SmallVector<int64_t, 4> loopTripCounts;
-};
-
-struct StaticProjectedLoopInfo {
-  mlir::BlockArgument iv;
-  int64_t lowerBound = 0;
-  int64_t step = 1;
-  int64_t tripCount = 1;
-};
-
-struct ProjectedTransferDescriptor {
-  ProjectedBatchInputKey inputKey;
-  mlir::Operation* extractOp = nullptr;
-  ProjectedFragmentLayout layout;
-  mlir::RankedTensorType payloadType;
-  llvm::SmallVector<llvm::SmallVector<int64_t, 4>, 16> fragmentOffsets;
-  llvm::SmallVector<llvm::SmallVector<int64_t, 16>, 4> fragmentOffsetsByDim;
-};
-
-struct ProjectedExtractReplacement {
-  mlir::Value payload;
-  ProjectedFragmentLayout layout;
-};
-
-struct ProjectedInputTransferFragment {
-  ProducerKey producer;
-  llvm::SmallVector<int64_t, 4> fragmentOffsets;
-  unsigned targetLane = 0;
-  unsigned ordinal = 0;
-  int64_t channelId = 0;
-  int32_t sourceCoreId = 0;
-  int32_t targetCoreId = 0;
-  bool sendEmitted = false;
-};
-
-struct ProjectedInputTransferPlan {
-  ProjectedBatchInputKey inputKey;
-  mlir::Operation* extractOp = nullptr;
-  ProjectedFragmentLayout layout;
-  llvm::SmallVector<ProjectedInputTransferFragment, 8> fragments;
-};
-
-struct PendingProjectedHostOutputFragment {
-  mlir::Value originalOutput;
-  ClassId sourceClass = 0;
-  ProducerKey producerKey;
-  unsigned publicationResultIndex = 0;
-  int64_t sourceFragmentOrdinal = 0;
-  int64_t sourceElementOffset = 0;
-  llvm::SmallVector<int64_t, 4> offsets;
-  llvm::SmallVector<int64_t, 4> sizes;
-  llvm::SmallVector<int64_t, 4> strides;
-  uint32_t sourceLane = 0;
-  mlir::Location loc;
-};
-
-struct AffineProjectedInputSliceMatch {
-  mlir::tensor::ExtractSliceOp extract;
-  mlir::RankedTensorType sourceType;
-  mlir::RankedTensorType fragmentType;
-  llvm::SmallVector<int64_t, 4> fragmentShape;
-  llvm::SmallVector<mlir::OpFoldResult, 4> offsets;
-  llvm::SmallVector<StaticProjectedLoopInfo, 4> loops;
-};
-
-unsigned getProjectedFragmentsPerLogicalSlot(llvm::ArrayRef<int64_t> loopTripCounts);
-mlir::LogicalResult verifyProjectedFragmentLayout(mlir::Operation* anchor, const ProjectedFragmentLayout& layout);
-mlir::FailureOr<mlir::RankedTensorType>
-getProjectedPayloadType(mlir::Operation* anchor, mlir::RankedTensorType fragmentType, unsigned payloadFragmentCount);
-llvm::SmallVector<llvm::SmallVector<int64_t, 16>, 4>
-buildProjectedFragmentOffsetsByDim(llvm::ArrayRef<llvm::SmallVector<int64_t, 4>> fragmentOffsets, size_t rank);
-mlir::LogicalResult verifyProjectedTransferDescriptor(mlir::Operation* anchor,
-                                                      const ProjectedTransferDescriptor& descriptor);
-mlir::LogicalResult verifyProjectedSendDescriptor(mlir::Operation* anchor,
-                                                  const ProjectedTransferDescriptor& descriptor,
-                                                  const MessageVector& messages);
-mlir::LogicalResult finalizeProjectedTransferDescriptor(mlir::Operation* anchor,
-                                                        ProjectedTransferDescriptor& descriptor);
-
-} // namespace onnx_mlir::spatial

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

@@ -106,10 +106,11 @@ static std::optional<uint32_t> getConstantExtractLane(tensor::ExtractSliceOp ext
 
 static std::optional<ProducerValueRef> getResultfulBatchProducerValueRef(SpatComputeBatch batch,
                                                                          const ComputeInstance* consumerInstance) {
-  if (!consumerInstance)
+  if (!consumerInstance || !isa<SpatComputeBatch>(consumerInstance->op))
-    return std::nullopt;
+    return ProducerValueRef {
-  if (!isa<SpatComputeBatch>(consumerInstance->op))
+      {batch.getOperation(), 0, static_cast<uint32_t>(batch.getLaneCount())},
-    return std::nullopt;
+      0
+    };
   if (consumerInstance->laneStart + consumerInstance->laneCount > static_cast<uint32_t>(batch.getLaneCount()))
     return std::nullopt;
   return ProducerValueRef {

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

@@ -0,0 +1,832 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/AsmState.h"
+#include "mlir/IR/IRMapping.h"
+#include "mlir/Pass/Pass.h"
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/raw_os_ostream.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include <fstream>
+#include <optional>
+#include <string>
+#include <vector>
+
+#include "ComputeInstanceUtils.hpp"
+#include "MergeSchedulingAnalysis.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
+#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Accelerators/PIM/Pass/PIMPasses.h"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace spatial {
+namespace {
+
+struct ProducerValueKey {
+  ComputeInstance instance;
+  size_t resultIndex = 0;
+
+  bool operator==(const ProducerValueKey& other) const {
+    return instance == other.instance && resultIndex == other.resultIndex;
+  }
+};
+
+struct GraphComputeBlockKey {
+  Operation* op = nullptr;
+  uint32_t laneStart = 0;
+  uint32_t laneCount = 1;
+
+  bool operator==(const GraphComputeBlockKey& other) const {
+    return op == other.op && laneStart == other.laneStart && laneCount == other.laneCount;
+  }
+};
+
+struct ScheduledMaterializationRecord {
+  Operation* scheduledOp = nullptr;
+  SmallVector<size_t> canonicalCpuClasses;
+  SmallVector<size_t> cpus;
+  SmallVector<GraphComputeBlockKey> computeKeys;
+  SmallVector<Block*> blocks;
+};
+
+} // namespace
+} // namespace spatial
+} // namespace onnx_mlir
+
+namespace llvm {
+template <>
+struct DenseMapInfo<onnx_mlir::spatial::ProducerValueKey> {
+  static onnx_mlir::spatial::ProducerValueKey getEmptyKey() {
+    return {DenseMapInfo<onnx_mlir::spatial::ComputeInstance>::getEmptyKey(), std::numeric_limits<size_t>::max()};
+  }
+  static onnx_mlir::spatial::ProducerValueKey getTombstoneKey() {
+    return {DenseMapInfo<onnx_mlir::spatial::ComputeInstance>::getTombstoneKey(), std::numeric_limits<size_t>::max()};
+  }
+  static unsigned getHashValue(const onnx_mlir::spatial::ProducerValueKey& key) {
+    return hash_combine(DenseMapInfo<onnx_mlir::spatial::ComputeInstance>::getHashValue(key.instance), key.resultIndex);
+  }
+  static bool isEqual(const onnx_mlir::spatial::ProducerValueKey& lhs,
+                      const onnx_mlir::spatial::ProducerValueKey& rhs) {
+    return lhs == rhs;
+  }
+};
+
+template <>
+struct DenseMapInfo<onnx_mlir::spatial::GraphComputeBlockKey> {
+  static onnx_mlir::spatial::GraphComputeBlockKey getEmptyKey() {
+    return {DenseMapInfo<mlir::Operation*>::getEmptyKey(), UINT32_MAX, UINT32_MAX};
+  }
+  static onnx_mlir::spatial::GraphComputeBlockKey getTombstoneKey() {
+    return {DenseMapInfo<mlir::Operation*>::getTombstoneKey(), UINT32_MAX, UINT32_MAX - 1};
+  }
+  static unsigned getHashValue(const onnx_mlir::spatial::GraphComputeBlockKey& key) {
+    return hash_combine(key.op, key.laneStart, key.laneCount);
+  }
+  static bool isEqual(const onnx_mlir::spatial::GraphComputeBlockKey& lhs,
+                      const onnx_mlir::spatial::GraphComputeBlockKey& rhs) {
+    return lhs == rhs;
+  }
+};
+} // namespace llvm
+
+namespace onnx_mlir {
+namespace spatial {
+namespace {
+
+struct ScalarClass {
+  size_t cpu = 0;
+  SmallVector<size_t> cpus;
+  SmallVector<ComputeInstance> instances;
+  SmallVector<Value> weights;
+  SmallVector<Value> inputs;
+  SmallVector<Type> resultTypes;
+};
+
+struct CompactCpuIndexMap {
+  DenseMap<size_t, size_t> rawToCompact;
+
+  size_t lookup(size_t rawCpu) const {
+    auto it = rawToCompact.find(rawCpu);
+    assert(it != rawToCompact.end() && "missing compact cpu index");
+    return it->second;
+  }
+};
+
+static ComputeInstance getWholeBatchInstance(SpatComputeBatch batch) {
+  return {batch.getOperation(), 0, static_cast<uint32_t>(batch.getLaneCount())};
+}
+
+static GraphComputeBlockKey getGraphComputeBlockKey(const ComputeInstance& instance) {
+  return {instance.op, instance.laneStart, instance.laneCount};
+}
+
+static CompactCpuIndexMap buildCompactCpuIndexMap(const MergeScheduleResult& schedule) {
+  SmallVector<size_t> rawCpus;
+  auto append = [&](size_t rawCpu) {
+    if (!llvm::is_contained(rawCpus, rawCpu))
+      rawCpus.push_back(rawCpu);
+  };
+  for (const auto& entry : schedule.computeToCpuMap)
+    append(entry.second);
+  for (const auto& entry : schedule.equivalentClass) {
+    append(entry.first);
+    for (size_t rawCpu : entry.second)
+      append(rawCpu);
+  }
+  for (const auto& entry : schedule.cpuToLastComputeMap)
+    append(entry.first);
+  llvm::sort(rawCpus);
+
+  CompactCpuIndexMap compactMap;
+  for (auto [index, rawCpu] : llvm::enumerate(rawCpus))
+    compactMap.rawToCompact[rawCpu] = index;
+  return compactMap;
+}
+
+static SmallVector<size_t> getEquivalentCpus(size_t cpu, const MergeScheduleResult& schedule) {
+  llvm::SmallDenseSet<size_t, 8> seen;
+  SmallVector<size_t> cpus;
+  auto append = [&](size_t value) {
+    if (seen.insert(value).second)
+      cpus.push_back(value);
+  };
+  append(cpu);
+  if (auto it = schedule.equivalentClass.find(cpu); it != schedule.equivalentClass.end())
+    for (size_t equivalentCpu : it->second)
+      append(equivalentCpu);
+  llvm::sort(cpus);
+  return cpus;
+}
+
+static size_t getCanonicalCpuClass(size_t cpu, const MergeScheduleResult& schedule) {
+  SmallVector<size_t> cpus = getEquivalentCpus(cpu, schedule);
+  return cpus.empty() ? cpu : cpus.front();
+}
+
+static Value getProducerOutput(const ProducerValueRef& ref) {
+  auto outputs = getComputeInstanceOutputValues(ref.instance);
+  assert(ref.resultIndex < outputs.size() && "producer result index out of range");
+  return outputs[ref.resultIndex];
+}
+
+static SmallVector<size_t> getCanonicalCpuClassesForBatch(SpatComputeBatch batch, const MergeScheduleResult& schedule) {
+  llvm::SmallDenseSet<size_t, 4> seen;
+  SmallVector<size_t> classes;
+  size_t chunkCount = getBatchChunkTargetCount(batch.getLaneCount());
+  for (size_t chunkIndex = 0; chunkIndex < chunkCount; ++chunkIndex) {
+    ComputeInstance chunk = getBatchChunkForIndex(batch, chunkIndex);
+    auto cpuIt = schedule.computeToCpuMap.find(chunk);
+    if (cpuIt == schedule.computeToCpuMap.end())
+      continue;
+    size_t cpuClass = getCanonicalCpuClass(cpuIt->second, schedule);
+    if (seen.insert(cpuClass).second)
+      classes.push_back(cpuClass);
+  }
+  llvm::sort(classes);
+  return classes;
+}
+
+static SmallVector<size_t> remapCpuList(ArrayRef<size_t> rawCpus, const CompactCpuIndexMap& compactCpuMap) {
+  SmallVector<size_t> compactCpus;
+  compactCpus.reserve(rawCpus.size());
+  for (size_t rawCpu : rawCpus)
+    compactCpus.push_back(compactCpuMap.lookup(rawCpu));
+  llvm::sort(compactCpus);
+  return compactCpus;
+}
+
+static std::string getInstanceName(const ComputeInstance& instance) {
+  return llvm::formatv("{0}[lanes={1}:{2}]",
+                       instance.op->getName().getStringRef(),
+                       instance.laneStart,
+                       instance.laneStart + instance.laneCount)
+    .str();
+}
+
+static void appendUnique(SmallVectorImpl<Value>& values, Value value) {
+  if (!llvm::is_contained(values, value))
+    values.push_back(value);
+}
+
+static Value getBlockOperand(Block& block, ValueRange operands, Value value, unsigned firstArgument = 0) {
+  auto it = llvm::find(operands, value);
+  assert(it != operands.end() && "missing scheduled operand");
+  return block.getArgument(firstArgument + std::distance(operands.begin(), it));
+}
+
+static void appendComputeBlockArguments(SmallVectorImpl<Type>& argTypes,
+                                        SmallVectorImpl<Location>& argLocs,
+                                        ValueRange weights,
+                                        ValueRange inputs,
+                                        ArrayRef<ProducerValueKey> carriedKeys,
+                                        Location loc) {
+  for (Value weight : weights)
+    argTypes.push_back(weight.getType());
+  for (Value input : inputs)
+    argTypes.push_back(input.getType());
+  for (ProducerValueKey key : carriedKeys) {
+    auto outputs = getComputeInstanceOutputValues(key.instance);
+    assert(key.resultIndex < outputs.size() && "missing carried result");
+    argTypes.push_back(outputs[key.resultIndex].getType());
+  }
+  argLocs.append(argTypes.size(), loc);
+}
+
+static Block* createScheduledComputeBlock(PatternRewriter& rewriter,
+                                          SpatScheduledCompute scheduled,
+                                          ArrayRef<ProducerValueKey> carriedKeys,
+                                          Location loc) {
+  SmallVector<Type> argTypes;
+  SmallVector<Location> argLocs;
+  appendComputeBlockArguments(argTypes, argLocs, scheduled.getWeights(), scheduled.getInputs(), carriedKeys, loc);
+  return rewriter.createBlock(&scheduled.getBody(), scheduled.getBody().end(), TypeRange(argTypes), argLocs);
+}
+
+static void appendBlockYieldBaseAndCarriedOperands(Block& block,
+                                                   unsigned baseArgCount,
+                                                   size_t carriedCount,
+                                                   SmallVectorImpl<Value>& operands) {
+  for (unsigned index = 0; index < baseArgCount; ++index)
+    operands.push_back(block.getArgument(index));
+  for (size_t index = 0; index < carriedCount; ++index)
+    operands.push_back(block.getArgument(baseArgCount + index));
+}
+
+static void createBlockYield(PatternRewriter& rewriter, Location loc, ValueRange outputs, Block* next = nullptr) {
+  OperationState state(loc, SpatBlockYieldOp::getOperationName());
+  state.addOperands(outputs);
+  if (next)
+    state.addSuccessors(next);
+  rewriter.create(state);
+}
+
+static std::string formatValueLabel(Value value, AsmState& asmState) {
+  std::string storage;
+  llvm::raw_string_ostream os(storage);
+  value.printAsOperand(os, asmState);
+  return storage;
+}
+
+static std::string formatBlockLabel(Block* block, AsmState& asmState) {
+  std::string storage;
+  llvm::raw_string_ostream os(storage);
+  block->printAsOperand(os, asmState);
+  return storage;
+}
+
+static std::string formatOperationLabel(Operation* op, AsmState& asmState) {
+  if (op->getNumResults() == 0)
+    return op->getName().getStringRef().str();
+  std::string storage;
+  llvm::raw_string_ostream os(storage);
+  llvm::interleaveComma(op->getResults(), os, [&](Value result) { os << formatValueLabel(result, asmState); });
+  return os.str();
+}
+
+static std::string formatGraphComputeBlockKey(const GraphComputeBlockKey& key, AsmState& asmState) {
+  return llvm::formatv("{0} {1} lanes [{2}:{3}]",
+                       formatOperationLabel(key.op, asmState),
+                       key.op->getName().getStringRef(),
+                       key.laneStart,
+                       key.laneStart + key.laneCount)
+    .str();
+}
+
+static size_t countPeftEquivalenceClasses(const MergeScheduleResult& schedule) {
+  llvm::SmallDenseSet<size_t, 16> classes;
+  for (const ComputeInstance& instance : schedule.dominanceOrderCompute) {
+    auto cpuIt = schedule.computeToCpuMap.find(instance);
+    if (cpuIt == schedule.computeToCpuMap.end())
+      continue;
+    classes.insert(getCanonicalCpuClass(cpuIt->second, schedule));
+  }
+  return classes.size();
+}
+
+static SmallVector<GraphComputeBlockKey> collectExpectedGraphComputeBlockKeys(func::FuncOp funcOp) {
+  SmallVector<GraphComputeBlockKey> keys;
+  for (Operation& op : funcOp.getOps()) {
+    if (auto compute = dyn_cast<SpatGraphCompute>(&op))
+      keys.push_back(getGraphComputeBlockKey({compute.getOperation(), 0, 1}));
+    else if (auto batch = dyn_cast<SpatGraphComputeBatch>(&op))
+      keys.push_back(getGraphComputeBlockKey(getWholeBatchInstance(batch)));
+  }
+  return keys;
+}
+
+static LogicalResult verifyMaterializedScheduleMapping(
+  func::FuncOp funcOp,
+  const MergeScheduleResult& schedule,
+  const DenseMap<GraphComputeBlockKey, Block*>& graphComputeToBlockMap,
+  ArrayRef<ScheduledMaterializationRecord> materializedSchedules) {
+  pim::CappedDiagnosticReporter diagnostics;
+  size_t expectedClassCount = countPeftEquivalenceClasses(schedule);
+  if (expectedClassCount != materializedSchedules.size()) {
+    diagnostics.report(funcOp.getOperation(), [&](Operation* illegalOp) {
+      illegalOp->emitOpError() << "phase-check expected " << expectedClassCount
+                               << " PEFT equivalence classes but materialized " << materializedSchedules.size()
+                               << " scheduled computes";
+    });
+  }
+
+  for (const ScheduledMaterializationRecord& record : materializedSchedules) {
+    if (record.canonicalCpuClasses.empty()) {
+      diagnostics.report(record.scheduledOp, [&](Operation* illegalOp) {
+        illegalOp->emitOpError("phase-check scheduled compute is missing an owning PEFT equivalence class");
+      });
+      continue;
+    }
+    if (record.canonicalCpuClasses.size() > 1) {
+      diagnostics.report(record.scheduledOp, [&](Operation* illegalOp) {
+        std::string classes;
+        llvm::raw_string_ostream os(classes);
+        llvm::interleaveComma(record.canonicalCpuClasses, os, [&](size_t cpuClass) { os << cpuClass; });
+        illegalOp->emitOpError("phase-check scheduled compute spans multiple PEFT equivalence classes")
+          << " (classes " << os.str() << ")";
+      });
+    }
+  }
+
+  for (GraphComputeBlockKey key : collectExpectedGraphComputeBlockKeys(funcOp)) {
+    if (graphComputeToBlockMap.count(key))
+      continue;
+    diagnostics.report(key.op, [&](Operation* illegalOp) {
+      illegalOp->emitOpError() << "phase-check graph compute is missing a scheduled MLIR block mapping for lanes ["
+                               << key.laneStart << ":" << (key.laneStart + key.laneCount) << "]";
+    });
+  }
+
+  for (const auto& entry : graphComputeToBlockMap) {
+    Block* block = entry.second;
+    if (!block || !isa<SpatScheduledCompute, SpatScheduledComputeBatch>(block->getParentOp())) {
+      diagnostics.report(entry.first.op, [&](Operation* illegalOp) {
+        illegalOp->emitOpError("phase-check graph compute block mapping does not target a scheduled compute block");
+      });
+    }
+  }
+
+  diagnostics.emitSuppressedSummary(funcOp, "scheduled Spatial materialization verification failed");
+  return success(!diagnostics.hasFailure());
+}
+
+static void dumpPeftMaterializationReport(func::FuncOp funcOp,
+                                          const MergeScheduleResult& schedule,
+                                          const CompactCpuIndexMap& compactCpuMap,
+                                          ArrayRef<ScheduledMaterializationRecord> materializedSchedules) {
+  std::fstream file = openDialectDumpFileWithExtension("spatial1_schedule_materialization", "/reports", "txt");
+  if (!file.is_open())
+    return;
+
+  llvm::raw_os_ostream os(file);
+  AsmState asmState(funcOp);
+
+  llvm::MapVector<size_t, SmallVector<size_t>> classToCpus;
+  DenseMap<size_t, SmallVector<ComputeInstance>> cpuToInstances;
+  for (const ComputeInstance& instance : schedule.dominanceOrderCompute) {
+    auto cpuIt = schedule.computeToCpuMap.find(instance);
+    if (cpuIt == schedule.computeToCpuMap.end())
+      continue;
+    size_t cpu = compactCpuMap.lookup(cpuIt->second);
+    size_t cpuClass = compactCpuMap.lookup(getCanonicalCpuClass(cpuIt->second, schedule));
+    if (!llvm::is_contained(classToCpus[cpuClass], cpu))
+      classToCpus[cpuClass].push_back(cpu);
+    cpuToInstances[cpu].push_back(instance);
+  }
+
+  os << "PEFT Schedule\n";
+  os << "=============\n";
+  os << "equivalence classes: " << classToCpus.size() << "\n\n";
+  for (auto& entry : classToCpus) {
+    size_t cpuClass = entry.first;
+    llvm::sort(entry.second);
+    os << "class " << cpuClass << " cpus [";
+    llvm::interleaveComma(entry.second, os, [&](size_t cpu) { os << cpu; });
+    os << "]\n";
+    for (size_t cpu : entry.second) {
+      os << "  cpu " << cpu << ":\n";
+      auto& instances = cpuToInstances[cpu];
+      llvm::sort(instances, [&](const ComputeInstance& lhs, const ComputeInstance& rhs) {
+        return schedule.computeToCpuSlotMap.lookup(lhs) < schedule.computeToCpuSlotMap.lookup(rhs);
+      });
+      for (const ComputeInstance& instance : instances) {
+        os << "    slot " << schedule.computeToCpuSlotMap.lookup(instance) << ": "
+           << formatGraphComputeBlockKey(getGraphComputeBlockKey(instance), asmState) << "\n";
+      }
+    }
+    os << "\n";
+  }
+
+  os << "Materialized Scheduled Computes\n";
+  os << "===============================\n";
+  os << "scheduled ops: " << materializedSchedules.size() << "\n\n";
+  for (const ScheduledMaterializationRecord& record : materializedSchedules) {
+    os << formatOperationLabel(record.scheduledOp, asmState) << " " << record.scheduledOp->getName().getStringRef();
+    os << " classes [";
+    llvm::interleaveComma(record.canonicalCpuClasses, os, [&](size_t cpuClass) { os << cpuClass; });
+    os << "] cpus [";
+    llvm::interleaveComma(record.cpus, os, [&](size_t cpu) { os << cpu; });
+    os << "]\n";
+    for (auto [key, block] : llvm::zip(record.computeKeys, record.blocks))
+      os << "  " << formatGraphComputeBlockKey(key, asmState) << " -> " << formatBlockLabel(block, asmState) << "\n";
+    os << "\n";
+  }
+}
+
+static LogicalResult createDeferredCommunication(PatternRewriter& rewriter,
+                                                 Location loc,
+                                                 Value consumerInput,
+                                                 Value source,
+                                                 const ProducerValueRef& producer,
+                                                 const ComputeInstance& consumer,
+                                                 const MergeScheduleResult& schedule,
+                                                 const CompactCpuIndexMap& compactCpuMap,
+                                                 uint64_t exchangeId,
+                                                 Value& result) {
+  auto sourceCpu = compactCpuMap.lookup(schedule.computeToCpuMap.lookup(producer.instance));
+  auto targetCpu = compactCpuMap.lookup(schedule.computeToCpuMap.lookup(consumer));
+  auto transfer = SpatDeferredCommunicationOp::create(
+    rewriter,
+    loc,
+    consumerInput.getType(),
+    ValueRange {source},
+    rewriter.getStringAttr(("x" + std::to_string(exchangeId)).c_str()),
+    rewriter.getStringAttr(getInstanceName(producer.instance)),
+    rewriter.getStringAttr(getInstanceName(consumer)),
+    rewriter.getI64IntegerAttr(static_cast<int64_t>(sourceCpu)),
+    rewriter.getI64IntegerAttr(static_cast<int64_t>(targetCpu)),
+    rewriter.getI64IntegerAttr(static_cast<int64_t>(sourceCpu)),
+    rewriter.getI64IntegerAttr(static_cast<int64_t>(targetCpu)),
+    rewriter.getI64IntegerAttr(producer.instance.laneStart),
+    rewriter.getI64IntegerAttr(consumer.laneStart),
+    rewriter.getStringAttr(consumerInput.getDefiningOp<tensor::ExtractSliceOp>() ? "projected" : "ordinary"),
+    rewriter.getI64IntegerAttr(static_cast<int64_t>(producer.resultIndex)),
+    rewriter.getDenseI64ArrayAttr({}),
+    rewriter.getI64IntegerAttr(-1));
+
+  Block* block = rewriter.createBlock(&transfer.getBody(), transfer.getBody().end(), source.getType(), loc);
+  rewriter.setInsertionPointToStart(block);
+  Value yielded = block->getArgument(0);
+  if (yielded.getType() != consumerInput.getType()) {
+    auto slice = consumerInput.getDefiningOp<tensor::ExtractSliceOp>();
+    if (!slice || slice.getSource().getType() != source.getType())
+      return transfer.emitOpError("cannot derive deferred communication payload type from recorded source");
+    IRMapping mapper;
+    mapper.map(slice.getSource(), yielded);
+    auto cloned = cast<tensor::ExtractSliceOp>(rewriter.clone(*slice, mapper));
+    yielded = cloned.getResult();
+  }
+  SpatYieldOp::create(rewriter, loc, yielded);
+  result = transfer.getOutput();
+  rewriter.setInsertionPointAfter(transfer);
+  return success();
+}
+
+static LogicalResult mapComputeInputs(PatternRewriter& rewriter,
+                                      Block& block,
+                                      SpatCompute compute,
+                                      const ComputeInstance& instance,
+                                      const ScalarClass& klass,
+                                      const MergeScheduleResult& schedule,
+                                      const CompactCpuIndexMap& compactCpuMap,
+                                      ValueRange scheduledWeights,
+                                      ValueRange scheduledInputs,
+                                      IRMapping& mapper,
+                                      uint64_t& exchangeId) {
+  for (auto [index, weight] : llvm::enumerate(compute.getWeights()))
+    mapper.map(*compute.getWeightArgument(index), getBlockOperand(block, scheduledWeights, weight));
+
+  for (auto [index, input] : llvm::enumerate(compute.getInputs())) {
+    Value mapped;
+    std::optional<ProducerValueRef> producer = getProducerValueRef(input, &instance);
+    if (producer) {
+      Value source = getProducerOutput(*producer);
+      if (failed(createDeferredCommunication(
+            rewriter, input.getLoc(), input, source, *producer, instance, schedule, compactCpuMap, exchangeId++, mapped)))
+        return failure();
+    }
+    else {
+      mapped = getBlockOperand(block, scheduledInputs, input, scheduledWeights.size());
+    }
+    mapper.map(*compute.getInputArgument(index), mapped);
+  }
+  return success();
+}
+
+static LogicalResult mapBatchInputs(PatternRewriter& rewriter,
+                                    Block& block,
+                                    SpatComputeBatch batch,
+                                    const MergeScheduleResult& schedule,
+                                    const CompactCpuIndexMap& compactCpuMap,
+                                    ValueRange scheduledWeights,
+                                    ValueRange scheduledInputs,
+                                    IRMapping& mapper,
+                                    uint64_t& exchangeId) {
+  mapper.map(*batch.getLaneArgument(), block.getArgument(0));
+  unsigned firstWeightArg = 1;
+  unsigned firstInputArg = firstWeightArg + scheduledWeights.size();
+  for (auto [index, weight] : llvm::enumerate(batch.getWeights()))
+    mapper.map(*batch.getWeightArgument(index), block.getArgument(firstWeightArg + index));
+
+  ComputeInstance consumer {batch.getOperation(), 0, static_cast<uint32_t>(batch.getLaneCount())};
+  for (auto [index, input] : llvm::enumerate(batch.getInputs())) {
+    Value mapped;
+    std::optional<ProducerValueRef> producer = getProducerValueRef(input, &consumer);
+    if (producer) {
+      Value source = getProducerOutput(*producer);
+      if (failed(createDeferredCommunication(
+            rewriter, input.getLoc(), input, source, *producer, consumer, schedule, compactCpuMap, exchangeId++, mapped)))
+        return failure();
+    }
+    else {
+      auto it = llvm::find(scheduledInputs, input);
+      if (it == scheduledInputs.end())
+        return batch.emitOpError("host batch input is missing from scheduled batch operands");
+      mapped = block.getArgument(firstInputArg + std::distance(scheduledInputs.begin(), it));
+    }
+    mapper.map(*batch.getInputArgument(index), mapped);
+  }
+
+  unsigned firstOutputArg = firstInputArg + scheduledInputs.size();
+  for (unsigned index = 0; index < batch.getNumResults(); ++index)
+    mapper.map(*batch.getOutputArgument(index), block.getArgument(firstOutputArg + index));
+  return success();
+}
+
+static void cloneComputeBody(PatternRewriter& rewriter,
+                             Block& source,
+                             IRMapping& mapper,
+                             SmallVectorImpl<Value>& yieldedValues) {
+  for (Operation& op : source.without_terminator())
+    rewriter.clone(op, mapper);
+  auto yield = cast<SpatYieldOp>(source.getTerminator());
+  for (Value output : yield.getOutputs())
+    yieldedValues.push_back(mapper.lookup(output));
+}
+
+struct MergeComputeNodesPass final : PassWrapper<MergeComputeNodesPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MergeComputeNodesPass)
+
+  StringRef getArgument() const override { return "pim-merge-compute-nodes"; }
+  StringRef getDescription() const override {
+    return "Materialize scheduled Spatial compute with deferred communication placeholders.";
+  }
+
+  void runOnOperation() override {
+    ModuleOp moduleOp = getOperation();
+    auto entryFunc = getPimEntryFunc(moduleOp);
+    if (failed(entryFunc)) {
+      moduleOp.emitError("failed to locate the PIM entry function during MergeComputeNodes");
+      signalPassFailure();
+      return;
+    }
+
+    func::FuncOp funcOp = *entryFunc;
+    MergeScheduleResult schedule = MergeSchedulingAnalysis(funcOp).getResult();
+    CompactCpuIndexMap compactCpuMap = buildCompactCpuIndexMap(schedule);
+    llvm::MapVector<size_t, ScalarClass> scalarClasses;
+    SmallVector<SpatComputeBatch> batchOps;
+    for (const ComputeInstance& instance : schedule.dominanceOrderCompute) {
+      if (isa<SpatCompute>(instance.op)) {
+        size_t cpu = schedule.computeToCpuMap.lookup(instance);
+        size_t cpuClass = getCanonicalCpuClass(cpu, schedule);
+        auto& klass = scalarClasses[cpuClass];
+        klass.cpu = cpuClass;
+        if (klass.cpus.empty())
+          klass.cpus = getEquivalentCpus(cpu, schedule);
+        klass.instances.push_back(instance);
+      }
+      else {
+        auto batch = cast<SpatComputeBatch>(instance.op);
+        if (!llvm::is_contained(batchOps, batch))
+          batchOps.push_back(batch);
+      }
+    }
+    for (auto& entry : scalarClasses) {
+      llvm::sort(entry.second.instances, [&](const ComputeInstance& lhs, const ComputeInstance& rhs) {
+        return schedule.computeToCpuSlotMap.lookup(lhs) < schedule.computeToCpuSlotMap.lookup(rhs);
+      });
+    }
+
+    PatternRewriter rewriter(moduleOp.getContext());
+    Operation* insertionPoint = funcOp.getBody().front().getTerminator();
+    DenseMap<ScalarClass*, SpatScheduledCompute> scalarOps;
+    DenseMap<Operation*, SpatScheduledComputeBatch> batchReplacements;
+    DenseMap<Operation*, size_t> scheduledOpToRecordIndex;
+    DenseMap<GraphComputeBlockKey, Block*> graphComputeToBlockMap;
+    std::vector<ScheduledMaterializationRecord> materializedSchedules;
+
+    for (auto& entry : scalarClasses) {
+      ScalarClass& klass = entry.second;
+      for (const ComputeInstance& instance : klass.instances) {
+        auto compute = cast<SpatCompute>(instance.op);
+        llvm::append_range(klass.resultTypes, compute.getResultTypes());
+        for (Value weight : compute.getWeights())
+          appendUnique(klass.weights, weight);
+        for (Value input : compute.getInputs()) {
+          std::optional<ProducerValueRef> producer = getProducerValueRef(input, &instance);
+          if (!producer)
+            appendUnique(klass.inputs, input);
+        }
+      }
+      rewriter.setInsertionPoint(insertionPoint);
+      auto scheduled = SpatScheduledCompute::create(
+        rewriter, klass.instances.front().op->getLoc(), TypeRange(klass.resultTypes), klass.weights, klass.inputs);
+      scheduled->setAttr(
+        kCoreIdAttrName, rewriter.getI32IntegerAttr(static_cast<int32_t>(compactCpuMap.lookup(klass.cpu))));
+      SmallVector<int64_t> equivalentCpus;
+      for (size_t cpu : klass.cpus)
+        equivalentCpus.push_back(static_cast<int64_t>(compactCpuMap.lookup(cpu)));
+      scheduled->setAttr("scheduled.equivalent_cpus", rewriter.getDenseI64ArrayAttr(equivalentCpus));
+      SmallVector<Attribute> sources;
+      SmallVector<int64_t> lanes;
+      for (const ComputeInstance& instance : klass.instances) {
+        sources.push_back(rewriter.getStringAttr(getInstanceName(instance)));
+        lanes.push_back(instance.laneStart);
+        lanes.push_back(instance.laneCount);
+      }
+      scheduled->setAttr("scheduled.sources", rewriter.getArrayAttr(sources));
+      scheduled->setAttr("scheduled.lane_ranges", rewriter.getDenseI64ArrayAttr(lanes));
+      scalarOps[&klass] = scheduled;
+      ScheduledMaterializationRecord record;
+      record.scheduledOp = scheduled.getOperation();
+      record.canonicalCpuClasses.push_back(compactCpuMap.lookup(klass.cpu));
+      record.cpus = remapCpuList(klass.cpus, compactCpuMap);
+      scheduledOpToRecordIndex[scheduled.getOperation()] = materializedSchedules.size();
+      materializedSchedules.push_back(std::move(record));
+    }
+
+    for (SpatComputeBatch batch : batchOps) {
+      SmallVector<Value> inputs;
+      ComputeInstance wholeBatch = getWholeBatchInstance(batch);
+      for (Value input : batch.getInputs())
+        if (!getProducerValueRef(input, &wholeBatch))
+          appendUnique(inputs, input);
+
+      rewriter.setInsertionPoint(insertionPoint);
+      auto scheduled = SpatScheduledComputeBatch::create(rewriter,
+                                                         batch.getLoc(),
+                                                         batch.getResultTypes(),
+                                                         rewriter.getI32IntegerAttr(batch.getLaneCount()),
+                                                         batch.getWeights(),
+                                                         inputs);
+      SmallVector<int32_t> coreIds;
+      uint32_t laneCount = batch.getLaneCount();
+      coreIds.reserve(laneCount);
+      for (uint32_t lane = 0; lane < laneCount; ++lane) {
+        ComputeInstance chunk = getBatchChunkForLane(batch, lane);
+        coreIds.push_back(static_cast<int32_t>(compactCpuMap.lookup(schedule.computeToCpuMap.lookup(chunk))));
+      }
+      scheduled->setAttr(kCoreIdsAttrName, rewriter.getDenseI32ArrayAttr(coreIds));
+      scheduled->setAttr("scheduled.sources", rewriter.getArrayAttr({rewriter.getStringAttr(getInstanceName(wholeBatch))}));
+      batchReplacements[batch.getOperation()] = scheduled;
+      ScheduledMaterializationRecord record;
+      record.scheduledOp = scheduled.getOperation();
+      record.canonicalCpuClasses = remapCpuList(getCanonicalCpuClassesForBatch(batch, schedule), compactCpuMap);
+      for (int32_t coreId : coreIds)
+        if (!llvm::is_contained(record.cpus, static_cast<size_t>(coreId)))
+          record.cpus.push_back(static_cast<size_t>(coreId));
+      scheduledOpToRecordIndex[scheduled.getOperation()] = materializedSchedules.size();
+      materializedSchedules.push_back(std::move(record));
+    }
+
+    uint64_t exchangeId = 0;
+    for (auto& entry : scalarClasses) {
+      ScalarClass& klass = entry.second;
+      SpatScheduledCompute scheduled = scalarOps.lookup(&klass);
+      SmallVector<ProducerValueKey> carriedKeys;
+      Block* block = nullptr;
+      for (auto [ordinal, instance] : llvm::enumerate(klass.instances)) {
+        auto compute = cast<SpatCompute>(instance.op);
+        if (!block)
+          block = createScheduledComputeBlock(rewriter, scheduled, carriedKeys, compute.getLoc());
+        GraphComputeBlockKey key = getGraphComputeBlockKey(instance);
+        graphComputeToBlockMap[key] = block;
+        auto& record = materializedSchedules[scheduledOpToRecordIndex.lookup(scheduled.getOperation())];
+        record.computeKeys.push_back(key);
+        record.blocks.push_back(block);
+        rewriter.setInsertionPointToStart(block);
+        IRMapping mapper;
+        if (failed(mapComputeInputs(rewriter,
+                                    *block,
+                                    compute,
+                                    instance,
+                                    klass,
+                                    schedule,
+                                    compactCpuMap,
+                                    scheduled.getWeights(),
+                                    scheduled.getInputs(),
+                                    mapper,
+                                    exchangeId))) {
+          signalPassFailure();
+          return;
+        }
+        SmallVector<Value> yieldedValues;
+        cloneComputeBody(rewriter, compute.getBody().front(), mapper, yieldedValues);
+        SmallVector<ProducerValueKey> currentKeys;
+        for (auto [index, value] : llvm::enumerate(yieldedValues)) {
+          ProducerValueKey key {instance, index};
+          currentKeys.push_back(key);
+        }
+
+        unsigned baseArgCount = scheduled.getWeights().size() + scheduled.getInputs().size();
+        SmallVector<Value> blockYieldOperands;
+        bool hasNextBlock = ordinal + 1 < klass.instances.size();
+        if (hasNextBlock) {
+          SmallVector<ProducerValueKey> nextCarriedKeys(carriedKeys);
+          llvm::append_range(nextCarriedKeys, currentKeys);
+          Block* nextBlock = createScheduledComputeBlock(rewriter, scheduled, nextCarriedKeys, compute.getLoc());
+          appendBlockYieldBaseAndCarriedOperands(*block, baseArgCount, carriedKeys.size(), blockYieldOperands);
+          llvm::append_range(blockYieldOperands, yieldedValues);
+          rewriter.setInsertionPointToEnd(block);
+          createBlockYield(rewriter, compute.getLoc(), blockYieldOperands, nextBlock);
+          carriedKeys = std::move(nextCarriedKeys);
+          block = nextBlock;
+        }
+        else {
+          for (size_t index = 0; index < carriedKeys.size(); ++index)
+            blockYieldOperands.push_back(block->getArgument(baseArgCount + index));
+          llvm::append_range(blockYieldOperands, yieldedValues);
+          rewriter.setInsertionPointToEnd(block);
+          createBlockYield(rewriter, compute.getLoc(), blockYieldOperands);
+        }
+      }
+    }
+
+    for (SpatComputeBatch batch : batchOps) {
+      SpatScheduledComputeBatch scheduled = batchReplacements.lookup(batch.getOperation());
+      SmallVector<Type> blockArgTypes {rewriter.getIndexType()};
+      SmallVector<Location> blockArgLocs {batch.getLoc()};
+      for (Value weight : scheduled.getWeights()) {
+        blockArgTypes.push_back(weight.getType());
+        blockArgLocs.push_back(weight.getLoc());
+      }
+      for (Value input : scheduled.getInputs()) {
+        blockArgTypes.push_back(input.getType());
+        blockArgLocs.push_back(input.getLoc());
+      }
+      for (Type resultType : scheduled.getResultTypes()) {
+        blockArgTypes.push_back(resultType);
+        blockArgLocs.push_back(batch.getLoc());
+      }
+      Block* block = rewriter.createBlock(
+        &scheduled.getBody(), scheduled.getBody().end(), TypeRange(blockArgTypes), blockArgLocs);
+      GraphComputeBlockKey key = getGraphComputeBlockKey(getWholeBatchInstance(batch));
+      graphComputeToBlockMap[key] = block;
+      auto& record = materializedSchedules[scheduledOpToRecordIndex.lookup(scheduled.getOperation())];
+      record.computeKeys.push_back(key);
+      record.blocks.push_back(block);
+      rewriter.setInsertionPointToStart(block);
+      IRMapping mapper;
+      if (failed(mapBatchInputs(rewriter,
+                                *block,
+                                batch,
+                                schedule,
+                                compactCpuMap,
+                                scheduled.getWeights(),
+                                scheduled.getInputs(),
+                                mapper,
+                                exchangeId))) {
+        signalPassFailure();
+        return;
+      }
+      for (Operation& op : batch.getBody().front())
+        rewriter.clone(op, mapper);
+    }
+
+    dumpPeftMaterializationReport(funcOp, schedule, compactCpuMap, materializedSchedules);
+    if (failed(verifyMaterializedScheduleMapping(funcOp, schedule, graphComputeToBlockMap, materializedSchedules))) {
+      moduleOp.emitError("scheduled Spatial materialization mapping verification failed");
+      signalPassFailure();
+      return;
+    }
+    if (failed(verifyScheduledSpatialInvariants(funcOp))) {
+      moduleOp.emitError("scheduled Spatial phase 1 verification failed");
+      signalPassFailure();
+      return;
+    }
+    dumpModule(moduleOp, "spatial1_scheduled_no_comm", /*assumeVerified=*/true);
+    moduleOp.emitError("MergeComputeNodes stopped after spatial1_scheduled_no_comm; "
+                       "Phase 2 communication realization is not implemented");
+    signalPassFailure();
+  }
+};
+
+} // namespace
+
+} // namespace spatial
+
+std::unique_ptr<Pass> createMergeComputeNodesPass() { return std::make_unique<spatial::MergeComputeNodesPass>(); }
+
+} // namespace onnx_mlir

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

@@ -1,728 +0,0 @@
-#include "SpatialDataflowCsvExporter.hpp"
-
-#include "mlir/Dialect/Affine/IR/AffineOps.h"
-#include "mlir/Dialect/Arith/IR/Arith.h"
-#include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/IR/AsmState.h"
-#include "mlir/IR/BuiltinAttributes.h"
-#include "mlir/IR/BuiltinOps.h"
-#include "mlir/IR/BuiltinTypes.h"
-#include "mlir/IR/Value.h"
-
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/SmallVector.h"
-#include "llvm/Support/Casting.h"
-#include "llvm/Support/ErrorHandling.h"
-#include "llvm/Support/raw_ostream.h"
-
-#include <cstdint>
-#include <fstream>
-#include <optional>
-#include <string>
-#include <utility>
-
-#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
-#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
-#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
-#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
-#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-
-using namespace mlir;
-
-namespace onnx_mlir {
-namespace spatial {
-
-namespace {
-
-struct TopLevelOpInfo {
-  Operation* op = nullptr;
-  size_t opId = 0;
-  bool isPost = false;
-  std::optional<int32_t> scalarCore;
-};
-
-struct ExpandedNodeInfo {
-  std::string id;
-  std::optional<int32_t> core;
-  std::optional<uint32_t> lane;
-};
-
-struct ChannelSendRecord {
-  std::string sourceId;
-  std::optional<uint32_t> sourceLane;
-};
-
-enum class LogicalNodeSelector {
-  Scalar,
-  Lane,
-  RangeRepresentative,
-};
-
-struct ResolvedProducer {
-  Operation* op = nullptr;
-  size_t resultIndex = 0;
-  LogicalNodeSelector selector = LogicalNodeSelector::Scalar;
-  uint32_t lane = 0;
-  uint32_t laneStart = 0;
-  uint32_t laneCount = 1;
-};
-
-struct EdgeSource {
-  std::string id;
-  std::optional<uint32_t> sourceLane;
-};
-
-std::string csvEscape(StringRef field) {
-  bool needsQuotes = field.contains(',') || field.contains('"') || field.contains('\n') || field.contains('\r');
-  if (!needsQuotes)
-    return field.str();
-
-  std::string escaped;
-  escaped.reserve(field.size() + 2);
-  escaped.push_back('"');
-  for (char ch : field) {
-    if (ch == '"')
-      escaped += "\"\"";
-    else
-      escaped.push_back(ch);
-  }
-  escaped.push_back('"');
-  return escaped;
-}
-
-void writeCsvRow(std::fstream& file, ArrayRef<std::string> fields) {
-  for (size_t i = 0; i < fields.size(); ++i) {
-    if (i != 0)
-      file << ",";
-    file << csvEscape(fields[i]);
-  }
-  file << "\n";
-}
-
-template <typename NumberT>
-std::string maybeNumber(std::optional<NumberT> value) {
-  if (!value)
-    return "";
-  return std::to_string(*value);
-}
-
-std::string stringifyType(Type type) {
-  std::string storage;
-  llvm::raw_string_ostream os(storage);
-  type.print(os);
-  return os.str();
-}
-
-std::string stringifyValueAsOperand(Value value, AsmState& asmState) {
-  std::string storage;
-  llvm::raw_string_ostream os(storage);
-  value.printAsOperand(os, asmState);
-  return os.str();
-}
-
-std::string stringifyResultSsaNames(Operation* op, AsmState* asmState) {
-  if (!asmState || op->getNumResults() == 0)
-    return "";
-
-  std::string storage;
-  llvm::raw_string_ostream os(storage);
-  llvm::interleave(
-    op->getResults(),
-    [&](Value result) { os << stringifyValueAsOperand(result, *asmState); },
-    [&]() { os << ";"; });
-  return os.str();
-}
-
-std::optional<uint64_t> getTypeSizeBytes(Type type) {
-  if (auto shapedType = dyn_cast<ShapedType>(type)) {
-    if (!shapedType.hasStaticShape() || !hasByteSizedElementType(shapedType.getElementType()))
-      return std::nullopt;
-    return static_cast<uint64_t>(getShapedTypeSizeInBytes(shapedType));
-  }
-
-  if (isa<IndexType>(type))
-    return static_cast<uint64_t>(getElementTypeSizeInBytes(type));
-  if (auto intType = dyn_cast<IntegerType>(type)) {
-    if (intType.getWidth() <= 0 || intType.getWidth() % 8 != 0)
-      return std::nullopt;
-    return static_cast<uint64_t>(getElementTypeSizeInBytes(type));
-  }
-  if (auto floatType = dyn_cast<FloatType>(type)) {
-    if (floatType.getWidth() <= 0 || floatType.getWidth() % 8 != 0)
-      return std::nullopt;
-    return static_cast<uint64_t>(getElementTypeSizeInBytes(type));
-  }
-  return std::nullopt;
-}
-
-std::string getScalarId(bool isPost, size_t opId) {
-  return (isPost ? "sc:" : "gc:") + std::to_string(opId);
-}
-
-std::string getBatchLaneId(bool isPost, size_t opId, uint32_t lane) {
-  return (isPost ? "scb:" : "gcb:") + std::to_string(opId) + ":" + std::to_string(lane);
-}
-
-template <typename ComputeOpTy, typename BatchOpTy>
-bool isTopLevelRelevantCompute(Operation& op) {
-  return isa<ComputeOpTy, BatchOpTy>(&op);
-}
-
-template <typename ComputeOpTy, typename BatchOpTy>
-FailureOr<TopLevelOpInfo> buildTopLevelOpInfo(Operation& op, bool isPost, size_t opId) {
-  TopLevelOpInfo info;
-  info.op = &op;
-  info.opId = opId;
-  info.isPost = isPost;
-
-  if constexpr (std::is_same_v<ComputeOpTy, SpatScheduledCompute>) {
-    if (auto compute = dyn_cast<ComputeOpTy>(&op)) {
-      auto coreId = getOptionalScheduledCoreId(compute, "spatial dataflow export core id");
-      if (failed(coreId))
-        return failure();
-      if (*coreId)
-        info.scalarCore = **coreId;
-    }
-  }
-
-  return info;
-}
-
-template <typename BatchOpTy>
-FailureOr<SmallVector<int32_t, 8>> getBatchLaneCoreIds(BatchOpTy batch) {
-  if constexpr (std::is_same_v<BatchOpTy, SpatScheduledComputeBatch>) {
-    auto coreIds = getOptionalScheduledBatchCoreIds(batch, "spatial dataflow export core ids");
-    if (failed(coreIds))
-      return failure();
-    if (!*coreIds)
-      return SmallVector<int32_t, 8> {};
-    return SmallVector<int32_t, 8>((**coreIds).begin(), (**coreIds).end());
-  }
-  return SmallVector<int32_t, 8> {};
-}
-
-std::string getExpandedNodeId(const DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes,
-                              Operation* op,
-                              uint32_t lane) {
-  auto it = expandedNodes.find({op, lane});
-  if (it == expandedNodes.end())
-    return "";
-  return it->second.id;
-}
-
-void addScalarNodeRow(std::fstream& nodesFile,
-                      DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes,
-                      const TopLevelOpInfo& info,
-                      AsmState* asmState = nullptr) {
-  std::string id = getScalarId(info.isPost, info.opId);
-  SmallVector<std::string, 5> row {id, std::to_string(info.opId), "", maybeNumber<int32_t>(info.scalarCore)};
-  if (asmState)
-    row.push_back(stringifyResultSsaNames(info.op, asmState));
-  writeCsvRow(nodesFile, row);
-  expandedNodes[{info.op, 0}] = {id, info.scalarCore, std::nullopt};
-}
-
-template <typename BatchOpTy>
-void addBatchNodeRows(std::fstream& nodesFile,
-                      DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes,
-                      const TopLevelOpInfo& info,
-                      BatchOpTy batch,
-                      ArrayRef<std::optional<int32_t>> laneCoreIds,
-                      AsmState* asmState = nullptr) {
-  for (uint32_t lane = 0; lane < static_cast<uint32_t>(batch.getLaneCount()); ++lane) {
-    std::string id = getBatchLaneId(info.isPost, info.opId, lane);
-    SmallVector<std::string, 5> row {id,
-                                     std::to_string(info.opId),
-                                     std::to_string(lane),
-                                     maybeNumber<int32_t>(laneCoreIds[lane])};
-    if (asmState)
-      row.push_back(stringifyResultSsaNames(info.op, asmState));
-    writeCsvRow(nodesFile, row);
-    expandedNodes[{info.op, lane}] = {id, laneCoreIds[lane], lane};
-  }
-}
-
-std::optional<int64_t> evaluateIndexLike(Value value, Value laneArg, uint32_t lane);
-
-std::optional<int64_t> evaluateIndexLike(Value value, Value laneArg, uint32_t lane) {
-  if (value == laneArg)
-    return static_cast<int64_t>(lane);
-
-  if (std::optional<int64_t> constant = matchConstantIndexValue(value))
-    return *constant;
-
-  if (auto constant = value.getDefiningOp<arith::ConstantOp>()) {
-    if (auto intAttr = dyn_cast<IntegerAttr>(constant.getValue()))
-      return intAttr.getInt();
-  }
-
-  if (auto extract = value.getDefiningOp<tensor::ExtractOp>()) {
-    auto constant = extract.getTensor().getDefiningOp<arith::ConstantOp>();
-    auto elements = constant ? dyn_cast<ElementsAttr>(constant.getValue()) : nullptr;
-    auto shapedType = elements ? dyn_cast<ShapedType>(elements.getType()) : nullptr;
-    if (!elements || !shapedType || shapedType.getRank() != 1 || extract.getIndices().size() != 1)
-      return std::nullopt;
-
-    std::optional<int64_t> index = evaluateIndexLike(extract.getIndices().front(), laneArg, lane);
-    if (!index || *index < 0 || *index >= static_cast<int64_t>(elements.getNumElements()))
-      return std::nullopt;
-
-    if (auto denseInts = dyn_cast<DenseIntElementsAttr>(elements))
-      return (*(denseInts.value_begin<APInt>() + *index)).getSExtValue();
-    return std::nullopt;
-  }
-
-  if (auto affineApply = value.getDefiningOp<affine::AffineApplyOp>())
-    if (FailureOr<int64_t> folded = evaluateAffineApply(
-          affineApply,
-          [&](Value operand) -> FailureOr<int64_t> {
-            if (std::optional<int64_t> resolved = evaluateIndexLike(operand, laneArg, lane))
-              return *resolved;
-            return failure();
-          });
-        succeeded(folded)) {
-      return *folded;
-    }
-
-  return std::nullopt;
-}
-
-SmallVector<int64_t, 8> collectPossibleIntValues(Value value, Value laneArg, uint32_t lane) {
-  if (std::optional<int64_t> exact = evaluateIndexLike(value, laneArg, lane))
-    return {*exact};
-
-  auto extract = value.getDefiningOp<tensor::ExtractOp>();
-  auto constant = extract ? extract.getTensor().getDefiningOp<arith::ConstantOp>() : nullptr;
-  auto elements = constant ? dyn_cast<ElementsAttr>(constant.getValue()) : nullptr;
-  if (!elements)
-    return {};
-
-  SmallVector<int64_t, 8> values;
-  if (auto denseInts = dyn_cast<DenseIntElementsAttr>(elements)) {
-    values.reserve(elements.getNumElements());
-    for (APInt element : denseInts.getValues<APInt>())
-      if (!llvm::is_contained(values, element.getSExtValue()))
-        values.push_back(element.getSExtValue());
-  }
-  return values;
-}
-
-template <typename BatchOpTy>
-std::optional<Value> getBatchLaneInput(BatchOpTy batch, uint32_t lane, unsigned inputIndex) {
-  if (batch.getNumResults() != 0)
-    return batch.getInputs()[inputIndex];
-
-  size_t laneCount = static_cast<size_t>(batch.getLaneCount());
-  if (laneCount == 0 || batch.getInputs().size() % laneCount != 0)
-    return std::nullopt;
-
-  size_t inputsPerLane = batch.getInputs().size() / laneCount;
-  size_t flatIndex = static_cast<size_t>(lane) * inputsPerLane + inputIndex;
-  if (flatIndex >= batch.getInputs().size())
-    return std::nullopt;
-  return batch.getInputs()[flatIndex];
-}
-
-template <typename BatchOpTy>
-unsigned getBatchLaneInputCount(BatchOpTy batch) {
-  if (batch.getNumResults() != 0)
-    return batch.getInputs().size();
-
-  size_t laneCount = static_cast<size_t>(batch.getLaneCount());
-  if (laneCount == 0 || batch.getInputs().size() % laneCount != 0)
-    return 0;
-  return static_cast<unsigned>(batch.getInputs().size() / laneCount);
-}
-
-template <typename ComputeOpTy, typename BatchOpTy>
-std::optional<ResolvedProducer> resolveProducerForValue(Value value, std::optional<uint32_t> consumerLane) {
-  Operation* op = value.getDefiningOp();
-  if (!op)
-    return std::nullopt;
-
-  while (auto extract = dyn_cast<tensor::ExtractSliceOp>(op)) {
-    Value source = extract.getSource();
-    Operation* sourceOp = source.getDefiningOp();
-    auto sourceBatch = dyn_cast_or_null<BatchOpTy>(sourceOp);
-    if (sourceBatch && sourceBatch.getNumResults() != 0) {
-      auto staticOffsets = extract.getStaticOffsets();
-      if (!staticOffsets.empty() && staticOffsets.front() != ShapedType::kDynamic) {
-        uint32_t lane = static_cast<uint32_t>(staticOffsets.front());
-        return ResolvedProducer {sourceOp, 0, LogicalNodeSelector::Lane, lane, lane, 1};
-      }
-      if (consumerLane)
-        return ResolvedProducer {sourceOp, 0, LogicalNodeSelector::Lane, *consumerLane, *consumerLane, 1};
-      return ResolvedProducer {
-        sourceOp, 0, LogicalNodeSelector::RangeRepresentative, 0, 0, static_cast<uint32_t>(sourceBatch.getLaneCount())
-      };
-    }
-    value = source;
-    op = sourceOp;
-    if (!op)
-      return std::nullopt;
-  }
-
-  if (auto compute = dyn_cast<ComputeOpTy>(op))
-    return ResolvedProducer {
-      compute.getOperation(), static_cast<size_t>(cast<OpResult>(value).getResultNumber()), LogicalNodeSelector::Scalar, 0, 0, 1
-    };
-
-  if (auto batch = dyn_cast<BatchOpTy>(op)) {
-    if (batch.getNumResults() != 0) {
-      if (consumerLane)
-        return ResolvedProducer {op, 0, LogicalNodeSelector::Lane, *consumerLane, *consumerLane, 1};
-      return ResolvedProducer {
-        op, 0, LogicalNodeSelector::RangeRepresentative, 0, 0, static_cast<uint32_t>(batch.getLaneCount())
-      };
-    }
-
-    uint32_t lane = static_cast<uint32_t>(cast<OpResult>(value).getResultNumber());
-    return ResolvedProducer {op, static_cast<size_t>(lane), LogicalNodeSelector::Lane, lane, lane, 1};
-  }
-
-  return std::nullopt;
-}
-
-SmallVector<EdgeSource, 8>
-resolveProducerSourcesForCsv(const ResolvedProducer& producer,
-                             const DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes) {
-  SmallVector<EdgeSource, 8> sources;
-
-  if (producer.selector == LogicalNodeSelector::Scalar) {
-    std::string id = getExpandedNodeId(expandedNodes, producer.op, 0);
-    if (!id.empty())
-      sources.push_back({id, std::nullopt});
-    return sources;
-  }
-
-  if (producer.selector == LogicalNodeSelector::Lane) {
-    std::string id = getExpandedNodeId(expandedNodes, producer.op, producer.lane);
-    if (!id.empty())
-      sources.push_back({id, producer.lane});
-    return sources;
-  }
-
-  for (uint32_t lane = producer.laneStart; lane < producer.laneStart + producer.laneCount; ++lane) {
-    std::string id = getExpandedNodeId(expandedNodes, producer.op, lane);
-    if (!id.empty())
-      sources.push_back({id, lane});
-  }
-  return sources;
-}
-
-void emitEdgeRow(std::fstream& edgesFile,
-                 StringRef sourceId,
-                 StringRef targetId,
-                 std::optional<uint64_t> byteSize,
-                 Type propagatedType,
-                 StringRef stage,
-                 std::optional<uint32_t> sourceLane,
-                 std::optional<uint32_t> targetLane,
-                 std::optional<int64_t> channelId) {
-  writeCsvRow(edgesFile,
-              {sourceId.str(),
-               targetId.str(),
-               maybeNumber<uint64_t>(byteSize),
-               stringifyType(propagatedType),
-               stage.str(),
-               maybeNumber<uint32_t>(sourceLane),
-               maybeNumber<uint32_t>(targetLane),
-               maybeNumber<int64_t>(channelId)});
-}
-
-template <typename ComputeOpTy, typename BatchOpTy>
-LogicalResult emitDataEdges(std::fstream& edgesFile,
-                            const DenseMap<Operation*, TopLevelOpInfo>& topLevelInfo,
-                            const DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes,
-                            StringRef stage) {
-  for (const auto& entry : topLevelInfo) {
-    Operation* op = entry.first;
-    const TopLevelOpInfo& info = entry.second;
-
-    if (auto compute = dyn_cast<ComputeOpTy>(op)) {
-      for (Value input : compute.getInputs()) {
-        if (isa_and_nonnull<SpatChannelReceiveOp>(input.getDefiningOp()))
-          continue;
-
-        auto producer = resolveProducerForValue<ComputeOpTy, BatchOpTy>(input, std::nullopt);
-        if (!producer)
-          continue;
-
-        SmallVector<EdgeSource, 8> sources = resolveProducerSourcesForCsv(*producer, expandedNodes);
-        std::optional<uint64_t> byteSize = getTypeSizeBytes(input.getType());
-        std::string targetId = getScalarId(info.isPost, info.opId);
-        for (const EdgeSource& source : sources)
-          emitEdgeRow(edgesFile, source.id, targetId, byteSize, input.getType(), stage, source.sourceLane, std::nullopt, std::nullopt);
-      }
-      continue;
-    }
-
-    auto batch = dyn_cast<BatchOpTy>(op);
-    if (!batch)
-      continue;
-
-    unsigned inputCount = getBatchLaneInputCount(batch);
-    for (uint32_t lane = 0; lane < static_cast<uint32_t>(batch.getLaneCount()); ++lane) {
-      std::string targetId = getBatchLaneId(info.isPost, info.opId, lane);
-      for (unsigned inputIndex = 0; inputIndex < inputCount; ++inputIndex) {
-        std::optional<Value> input = getBatchLaneInput(batch, lane, inputIndex);
-        if (!input || isa_and_nonnull<SpatChannelReceiveOp>((*input).getDefiningOp()))
-          continue;
-
-        auto producer = resolveProducerForValue<ComputeOpTy, BatchOpTy>(*input, lane);
-        if (!producer)
-          continue;
-
-        SmallVector<EdgeSource, 8> sources = resolveProducerSourcesForCsv(*producer, expandedNodes);
-        std::optional<uint64_t> byteSize = getTypeSizeBytes((*input).getType());
-        for (const EdgeSource& source : sources)
-          emitEdgeRow(edgesFile, source.id, targetId, byteSize, (*input).getType(), stage, source.sourceLane, lane, std::nullopt);
-      }
-    }
-  }
-
-  return success();
-}
-
-template <typename BatchOpTy>
-void collectChannelSends(DenseMap<int64_t, SmallVector<ChannelSendRecord, 4>>& sendsByChannelId,
-                         const DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes,
-                         BatchOpTy batch) {
-  std::optional<BlockArgument> laneArg = batch.getLaneArgument();
-  if (!laneArg)
-    return;
-
-  for (uint32_t lane = 0; lane < static_cast<uint32_t>(batch.getLaneCount()); ++lane) {
-    std::string sourceId = getExpandedNodeId(expandedNodes, batch.getOperation(), lane);
-    if (sourceId.empty())
-      continue;
-    batch.getBody().walk([&](SpatChannelSendOp send) {
-      std::optional<int64_t> channelId = evaluateIndexLike(send.getChannelId(), *laneArg, lane);
-      if (!channelId)
-        return;
-      sendsByChannelId[*channelId].push_back({sourceId, lane});
-    });
-  }
-}
-
-void collectChannelSends(DenseMap<int64_t, SmallVector<ChannelSendRecord, 4>>& sendsByChannelId,
-                         const DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes,
-                         SpatScheduledCompute compute) {
-  std::string sourceId = getExpandedNodeId(expandedNodes, compute.getOperation(), 0);
-  if (sourceId.empty())
-    return;
-  compute.getBody().walk([&](SpatChannelSendOp send) {
-    std::optional<int64_t> channelId = evaluateIndexLike(send.getChannelId(), Value(), 0);
-    if (!channelId)
-      return;
-    sendsByChannelId[*channelId].push_back({sourceId, std::nullopt});
-  });
-}
-
-DenseMap<int32_t, SmallVector<ChannelSendRecord, 4>>
-buildNodesByCore(const DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes) {
-  DenseMap<int32_t, SmallVector<ChannelSendRecord, 4>> nodesByCore;
-  for (const auto& entry : expandedNodes) {
-    const ExpandedNodeInfo& node = entry.second;
-    if (!node.core)
-      continue;
-    nodesByCore[*node.core].push_back({node.id, node.lane});
-  }
-  return nodesByCore;
-}
-
-template <typename ComputeOpTy, typename BatchOpTy, typename ResolveChannelSourcesFn>
-LogicalResult emitExplicitChannelEdges(std::fstream& edgesFile,
-                                       const DenseMap<Operation*, TopLevelOpInfo>& topLevelInfo,
-                                       ResolveChannelSourcesFn&& resolveChannelSources,
-                                       StringRef stage) {
-  for (const auto& entry : topLevelInfo) {
-    Operation* op = entry.first;
-    const TopLevelOpInfo& info = entry.second;
-
-    if (auto compute = dyn_cast<ComputeOpTy>(op)) {
-      compute.getBody().walk([&](SpatChannelReceiveOp receive) {
-        SmallVector<ChannelSendRecord, 4> sources = resolveChannelSources(receive, 0);
-        if (sources.empty())
-          return;
-        std::optional<int64_t> channelId = evaluateIndexLike(receive.getChannelId(), Value(), 0);
-        std::string targetId = getScalarId(info.isPost, info.opId);
-        std::optional<uint64_t> byteSize = getTypeSizeBytes(receive.getType());
-        for (const ChannelSendRecord& source : sources)
-          emitEdgeRow(edgesFile, source.sourceId, targetId, byteSize, receive.getType(), stage, source.sourceLane, std::nullopt, channelId);
-      });
-      continue;
-    }
-
-    auto batch = dyn_cast<BatchOpTy>(op);
-    if (!batch)
-      continue;
-    auto laneArg = batch.getLaneArgument();
-    if (!laneArg)
-      continue;
-    for (uint32_t lane = 0; lane < static_cast<uint32_t>(batch.getLaneCount()); ++lane) {
-      std::string targetId = getBatchLaneId(info.isPost, info.opId, lane);
-      batch.getBody().walk([&](SpatChannelReceiveOp receive) {
-        SmallVector<ChannelSendRecord, 4> sources = resolveChannelSources(receive, lane);
-        if (sources.empty())
-          return;
-        std::optional<int64_t> channelId = evaluateIndexLike(receive.getChannelId(), *laneArg, lane);
-        std::optional<uint64_t> byteSize = getTypeSizeBytes(receive.getType());
-        for (const ChannelSendRecord& source : sources)
-          emitEdgeRow(edgesFile, source.sourceId, targetId, byteSize, receive.getType(), stage, source.sourceLane, lane, channelId);
-      });
-    }
-  }
-
-  return success();
-}
-
-LogicalResult exportStagePre(func::FuncOp func) {
-  std::fstream nodesFile = openDialectDumpFileWithExtension("spatial1_graph.nodes", "/reports", "csv");
-  std::fstream edgesFile = openDialectDumpFileWithExtension("spatial1_graph.edges","/reports", "csv");
-  if (!nodesFile.is_open() || !edgesFile.is_open())
-    return success();
-
-  writeCsvRow(nodesFile, {"Id", "op_id", "lane", "core", "ssa_name"});
-  writeCsvRow(edgesFile, {"Source", "Target", "Weight", "Type", "stage", "source_lane", "target_lane", "channel_id"});
-
-  Operation* asmRoot = func.getOperation();
-  if (auto moduleOp = func->getParentOfType<ModuleOp>())
-    asmRoot = moduleOp.getOperation();
-  OpPrintingFlags flags;
-  flags.elideLargeElementsAttrs().enableDebugInfo(true, false);
-  AsmState asmState(asmRoot, flags);
-
-  DenseMap<Operation*, TopLevelOpInfo> topLevelInfo;
-  DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo> expandedNodes;
-
-  size_t opId = 0;
-  for (Operation& op : func.getBody().front()) {
-    if (!isTopLevelRelevantCompute<SpatGraphCompute, SpatGraphComputeBatch>(op))
-      continue;
-    FailureOr<TopLevelOpInfo> info = buildTopLevelOpInfo<SpatGraphCompute, SpatGraphComputeBatch>(op, false, opId++);
-    if (failed(info))
-      return failure();
-    topLevelInfo[&op] = *info;
-
-    if (auto compute = dyn_cast<SpatGraphCompute>(&op)) {
-      addScalarNodeRow(nodesFile, expandedNodes, *info, &asmState);
-      continue;
-    }
-
-    auto batch = cast<SpatGraphComputeBatch>(&op);
-    SmallVector<std::optional<int32_t>, 8> laneCoreIds(batch.getLaneCount());
-    addBatchNodeRows(nodesFile, expandedNodes, *info, batch, laneCoreIds, &asmState);
-  }
-
-  return emitDataEdges<SpatGraphCompute, SpatGraphComputeBatch>(edgesFile, topLevelInfo, expandedNodes, "pre");
-}
-
-LogicalResult exportStagePost(func::FuncOp func) {
-  std::fstream nodesFile = openDialectDumpFileWithExtension("spatial2_merged.nodes", "/reports", "csv");
-  std::fstream edgesFile = openDialectDumpFileWithExtension("spatial2_merged.edges", "/reports", "csv");
-  if (!nodesFile.is_open() || !edgesFile.is_open())
-    return success();
-
-  writeCsvRow(nodesFile, {"Id", "op_id", "lane", "core"});
-  writeCsvRow(edgesFile, {"Source", "Target", "Weight", "Type", "stage", "source_lane", "target_lane", "channel_id"});
-
-  DenseMap<Operation*, TopLevelOpInfo> topLevelInfo;
-  DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo> expandedNodes;
-
-  size_t opId = 0;
-  for (Operation& op : func.getBody().front()) {
-    if (!isTopLevelRelevantCompute<SpatScheduledCompute, SpatScheduledComputeBatch>(op))
-      continue;
-    FailureOr<TopLevelOpInfo> info = buildTopLevelOpInfo<SpatScheduledCompute, SpatScheduledComputeBatch>(op, true, opId++);
-    if (failed(info))
-      return failure();
-    topLevelInfo[&op] = *info;
-
-    if (isa<SpatScheduledCompute>(&op)) {
-      addScalarNodeRow(nodesFile, expandedNodes, *info);
-      continue;
-    }
-
-    auto batch = cast<SpatScheduledComputeBatch>(&op);
-    auto coreIds = getBatchLaneCoreIds(batch);
-    if (failed(coreIds))
-      return failure();
-    SmallVector<std::optional<int32_t>, 8> laneCoreIds(batch.getLaneCount());
-    for (uint32_t lane = 0; lane < static_cast<uint32_t>(batch.getLaneCount()); ++lane)
-      if (lane < coreIds->size())
-        laneCoreIds[lane] = (*coreIds)[lane];
-    addBatchNodeRows(nodesFile, expandedNodes, *info, batch, laneCoreIds);
-  }
-
-  if (failed(emitDataEdges<SpatScheduledCompute, SpatScheduledComputeBatch>(edgesFile, topLevelInfo, expandedNodes, "post")))
-    return failure();
-
-  DenseMap<int64_t, SmallVector<ChannelSendRecord, 4>> sendsByChannelId;
-  for (const auto& entry : topLevelInfo) {
-    Operation* op = entry.first;
-    if (auto compute = dyn_cast<SpatScheduledCompute>(op))
-      collectChannelSends(sendsByChannelId, expandedNodes, compute);
-    else if (auto batch = dyn_cast<SpatScheduledComputeBatch>(op))
-      collectChannelSends(sendsByChannelId, expandedNodes, batch);
-  }
-
-  DenseMap<int32_t, SmallVector<ChannelSendRecord, 4>> nodesByCore = buildNodesByCore(expandedNodes);
-  auto resolveChannelSources = [&](SpatChannelReceiveOp receive, uint32_t lane) {
-    SmallVector<ChannelSendRecord, 4> sources;
-
-    Value laneArg;
-    if (auto owner = receive->getParentOfType<SpatScheduledComputeBatch>())
-      if (auto maybeLaneArg = owner.getLaneArgument())
-        laneArg = *maybeLaneArg;
-
-    if (std::optional<int64_t> channelId = evaluateIndexLike(receive.getChannelId(), laneArg, lane)) {
-      if (auto it = sendsByChannelId.find(*channelId); it != sendsByChannelId.end())
-        return it->second;
-    }
-
-    for (int64_t sourceCore : collectPossibleIntValues(receive.getSourceCoreId(), laneArg, lane)) {
-      auto it = nodesByCore.find(static_cast<int32_t>(sourceCore));
-      if (it == nodesByCore.end())
-        continue;
-      llvm::append_range(sources, it->second);
-    }
-    return sources;
-  };
-
-  return emitExplicitChannelEdges<SpatScheduledCompute, SpatScheduledComputeBatch>(
-    edgesFile, topLevelInfo, resolveChannelSources, "post");
-}
-
-} // namespace
-
-SpatialDataflowExportStage getSpatialDataflowExportStage() {
-  switch (pimExportSpatialDataflow.getValue()) {
-  case SpatialDataflowExportNone: return SpatialDataflowExportStage::None;
-  case SpatialDataflowExportPre: return SpatialDataflowExportStage::Pre;
-  case SpatialDataflowExportPost: return SpatialDataflowExportStage::Post;
-  case SpatialDataflowExportBoth: return SpatialDataflowExportStage::Both;
-  }
-  llvm_unreachable("unknown spatial dataflow export mode");
-}
-
-bool shouldExportSpatialDataflowStage(SpatialDataflowExportStage mode, SpatialDataflowExportStage stage) {
-  switch (mode) {
-  case SpatialDataflowExportStage::None: return false;
-  case SpatialDataflowExportStage::Pre: return stage == SpatialDataflowExportStage::Pre;
-  case SpatialDataflowExportStage::Post: return stage == SpatialDataflowExportStage::Post;
-  case SpatialDataflowExportStage::Both:
-    return stage == SpatialDataflowExportStage::Pre || stage == SpatialDataflowExportStage::Post;
-  }
-  return false;
-}
-
-LogicalResult exportSpatialDataflowCsvPre(func::FuncOp func) { return exportStagePre(func); }
-
-LogicalResult exportSpatialDataflowCsvPost(func::FuncOp func) { return exportStagePost(func); }
-
-} // namespace spatial
-} // namespace onnx_mlir

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp

@@ -1,25 +0,0 @@
-#pragma once
-
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/Support/LogicalResult.h"
-
-
-namespace onnx_mlir {
-namespace spatial {
-
-enum class SpatialDataflowExportStage {
-  None,
-  Pre,
-  Post,
-  Both,
-};
-
-SpatialDataflowExportStage getSpatialDataflowExportStage();
-
-mlir::LogicalResult exportSpatialDataflowCsvPre(mlir::func::FuncOp func);
-mlir::LogicalResult exportSpatialDataflowCsvPost(mlir::func::FuncOp func);
-
-bool shouldExportSpatialDataflowStage(SpatialDataflowExportStage mode, SpatialDataflowExportStage stage);
-
-} // namespace spatial
-} // namespace onnx_mlir