CommunicationPlan

src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Conv.cpp

@@ -2716,6 +2716,181 @@ static FailureOr<Value> createNchwRowStripConvPatchRow(Value paddedWindow,
     .getResult();
 }
 
+static FailureOr<Value> createPaddedConvOutputRow(Value patchRow,
+                                                  const ConvLoweringState& state,
+                                                  Value paddedWeights,
+                                                  Value paddedBias,
+                                                  int64_t paddedK,
+                                                  int64_t numKSlices,
+                                                  int64_t xbarDim,
+                                                  PatternRewriter& rewriter,
+                                                  Location loc) {
+  const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
+  auto elementType = state.outType.getElementType();
+  auto rowType = RankedTensorType::get({1, state.numChannelsOut}, elementType);
+  auto paddedRowType = RankedTensorType::get({1, xbarDim}, elementType);
+  auto paddedPatchRowType = RankedTensorType::get({1, paddedK}, elementType);
+  auto paddedWeightTileType = RankedTensorType::get({xbarDim, xbarDim}, state.wType.getElementType());
+
+  Value paddedPatchRow = patchRow;
+  if (patchSize != paddedK)
+    paddedPatchRow = createZeroPaddedTensor(
+      paddedPatchRow, paddedPatchRowType, {0, 0}, {0, paddedK - patchSize}, rewriter, loc);
+
+  Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+  Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
+  Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
+  Value cNumKSlices = getOrCreateIndexConstant(rewriter, anchorOp, numKSlices);
+  Value cXbar = getOrCreateIndexConstant(rewriter, anchorOp, xbarDim);
+  auto createPiece = [&](Value kSlice, Location pieceLoc) -> Value {
+    Value kOffset = arith::MulIOp::create(rewriter, pieceLoc, kSlice, cXbar);
+    SmallVector<OpFoldResult> aOffsets {rewriter.getIndexAttr(0), kOffset};
+    SmallVector<OpFoldResult> aSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)};
+    Value aTile = extractStaticSliceOrIdentity(
+      rewriter, pieceLoc, paddedPatchRow, paddedRowType, aOffsets, aSizes, getUnitStrides(rewriter, 2));
+    SmallVector<OpFoldResult> bOffsets {kOffset, rewriter.getIndexAttr(0)};
+    SmallVector<OpFoldResult> bSizes {rewriter.getIndexAttr(xbarDim), rewriter.getIndexAttr(xbarDim)};
+    Value bTile = extractStaticSliceOrIdentity(
+      rewriter, pieceLoc, paddedWeights, paddedWeightTileType, bOffsets, bSizes, getUnitStrides(rewriter, 2));
+    return spatial::SpatVMMOp::create(rewriter, pieceLoc, paddedRowType, bTile, aTile).getResult();
+  };
+
+  Value rowResult = createPiece(c0, loc);
+  if (numKSlices > 1) {
+    auto kLoop = buildNormalizedScfFor(
+      rewriter,
+      loc,
+      c1,
+      cNumKSlices,
+      c1,
+      ValueRange {rowResult},
+      [&](OpBuilder&, Location reduceLoc, Value kSlice, ValueRange reduceIterArgs, SmallVectorImpl<Value>& reduceYielded) {
+        Value piece = createPiece(kSlice, reduceLoc);
+        reduceYielded.push_back(
+          spatial::SpatVAddOp::create(rewriter, reduceLoc, paddedRowType, reduceIterArgs.front(), piece).getResult());
+        return success();
+      });
+    if (failed(kLoop))
+      return failure();
+    rowResult = kLoop->results.front();
+  }
+
+  if (paddedBias)
+    rowResult = spatial::SpatVAddOp::create(rewriter, loc, paddedRowType, rowResult, paddedBias).getResult();
+  if (state.numChannelsOut == xbarDim)
+    return rowResult;
+
+  SmallVector<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
+  SmallVector<OpFoldResult> outputSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut)};
+  return tensor::ExtractSliceOp::create(
+           rewriter, loc, rowType, rowResult, outputOffsets, outputSizes, getUnitStrides(rewriter, 2))
+    .getResult();
+}
+
+static FailureOr<Value>
+createRowStripConvOutputFromDenseInput(const ConvLoweringState& state, PatternRewriter& rewriter, Location loc) {
+  ConvGeometry geometry = buildConvGeometry(state);
+  if (state.group != 1 || state.batchSize != 1 || geometry.c > geometry.xbarSize)
+    return failure();
+
+  auto weightDenseAttr = getHostConstDenseElementsAttr(state.w);
+  if (!weightDenseAttr)
+    return failure();
+  if (state.hasBias && !isSupportedBiasAddValue(state.b, state.outType))
+    return failure();
+
+  const int64_t xbarDim = geometry.xbarSize;
+  const int64_t patchSize = state.numChannelsIn * state.wHeight * state.wWidth;
+  const int64_t numKSlices = ceilIntegerDivide(patchSize, xbarDim);
+  const int64_t paddedK = numKSlices * xbarDim;
+  auto elementType = state.outType.getElementType();
+  auto fragmentType = getRowStripFragmentType(state.outType);
+  auto outputPixelType = RankedTensorType::get({1, state.numChannelsOut, 1, 1}, elementType);
+  auto patchType = RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.wWidth}, state.xType.getElementType());
+  auto patchRowType = RankedTensorType::get({1, patchSize}, state.xType.getElementType());
+  auto outputStorageType = getRowStripStorageType(state.outType);
+
+  PreparedConvInput preparedInput = standard::prepareInputForIm2Col(state, rewriter, loc);
+  Value paddedWeights = standard::createPaddedInputKTiledWeightConstant(weightDenseAttr, state, paddedK, xbarDim, rewriter);
+  FailureOr<Value> paddedBias = failure();
+  if (state.hasBias)
+    paddedBias = createPaddedBiasRowConstant(state, xbarDim, rewriter);
+  if (state.hasBias && failed(paddedBias))
+    return failure();
+
+  auto batchOp = createSpatComputeBatch(
+    rewriter,
+    loc,
+    TypeRange {outputStorageType},
+    state.outHeight,
+    ValueRange {paddedWeights},
+    state.hasBias ? ValueRange {preparedInput.value, *paddedBias} : ValueRange {preparedInput.value},
+    [&](detail::SpatComputeBatchBodyArgs args) {
+      Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+      Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
+      Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
+      Value cOutWidth = getOrCreateIndexConstant(rewriter, anchorOp, state.outWidth);
+      Value inputHeightOffset = affineMulConst(rewriter, loc, args.lane, state.strideHeight, anchorOp);
+      Value fragmentInit = tensor::EmptyOp::create(rewriter, loc, fragmentType.getShape(), elementType);
+      auto widthLoop = buildNormalizedScfFor(
+        rewriter,
+        loc,
+        c0,
+        cOutWidth,
+        c1,
+        ValueRange {fragmentInit},
+        [&](OpBuilder&, Location widthLoc, Value widthIndex, ValueRange widthIterArgs, SmallVectorImpl<Value>& widthYielded) {
+          Value inputWidthOffset = affineMulConst(rewriter, widthLoc, widthIndex, state.strideWidth, anchorOp);
+          Value patch = createConvInputPatch(args.inputs.front(),
+                                             patchType,
+                                             c0,
+                                             c0,
+                                             inputHeightOffset,
+                                             inputWidthOffset,
+                                             state.dilationHeight,
+                                             state.dilationWidth,
+                                             rewriter,
+                                             widthLoc);
+          Value patchRow = tensor::CollapseShapeOp::create(
+            rewriter, widthLoc, patchRowType, patch, SmallVector<ReassociationIndices> {{0}, {1, 2, 3}});
+          FailureOr<Value> outputRow = createPaddedConvOutputRow(patchRow,
+                                                                 state,
+                                                                 args.weights.front(),
+                                                                 state.hasBias ? args.inputs[1] : Value(),
+                                                                 paddedK,
+                                                                 numKSlices,
+                                                                 xbarDim,
+                                                                 rewriter,
+                                                                 widthLoc);
+          if (failed(outputRow))
+            return failure();
+
+          Value outputFragment = tensor::ExpandShapeOp::create(rewriter,
+                                                               widthLoc,
+                                                               outputPixelType,
+                                                               *outputRow,
+                                                               SmallVector<ReassociationIndices> {{0}, {1, 2, 3}});
+          SmallVector<OpFoldResult> rowOffsets {
+            rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex};
+          SmallVector<OpFoldResult> rowSizes {
+            rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut), rewriter.getIndexAttr(1),
+            rewriter.getIndexAttr(1)};
+          Value nextFragment = tensor::InsertSliceOp::create(
+            rewriter, widthLoc, outputFragment, widthIterArgs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 4));
+          widthYielded.push_back(nextFragment);
+          return success();
+        });
+      if (failed(widthLoop))
+        return failure();
+
+      insertRowStripFragment(widthLoop->results.front(), args.outputs.front(), state.outType, args.lane, rewriter, loc);
+      return success();
+    });
+  if (failed(batchOp))
+    return failure();
+  return batchOp->getResult(0);
+}
+
 static FailureOr<Value> createConvOutputFromNchwRowStripFragments(Value rowStripStorage,
                                                                   const ConvLoweringState& state,
                                                                   PatternRewriter& rewriter,
@@ -2734,11 +2909,7 @@ static FailureOr<Value> createConvOutputFromNchwRowStripFragments(Value rowStrip
   const int64_t numKSlices = ceilIntegerDivide(patchSize, xbarDim);
   const int64_t paddedK = numKSlices * xbarDim;
   auto elementType = state.outType.getElementType();
-  auto rowType = RankedTensorType::get({1, state.numChannelsOut}, state.outType.getElementType());
   auto outputPixelType = RankedTensorType::get({1, state.numChannelsOut, 1, 1}, elementType);
-  auto paddedRowType = RankedTensorType::get({1, xbarDim}, state.outType.getElementType());
-  auto paddedPatchRowType = RankedTensorType::get({1, paddedK}, elementType, inputType.getEncoding());
-  auto paddedWeightTileType = RankedTensorType::get({xbarDim, xbarDim}, state.wType.getElementType());
   auto outputStorageType = getRowStripStorageType(state.outType);
   auto weightDenseAttr = getHostConstDenseElementsAttr(state.w);
   if (!weightDenseAttr)
@@ -2761,9 +2932,7 @@ static FailureOr<Value> createConvOutputFromNchwRowStripFragments(Value rowStrip
       Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
       Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
       Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
-      Value cNumKSlices = getOrCreateIndexConstant(rewriter, anchorOp, numKSlices);
       Value cOutWidth = getOrCreateIndexConstant(rewriter, anchorOp, state.outWidth);
-      Value cXbar = getOrCreateIndexConstant(rewriter, anchorOp, xbarDim);
       auto fragmentType = getRowStripFragmentType(state.outType);
       FailureOr<Value> inputWindow = createNchwRowStripConvWindow(args.inputs.front(), state, args.lane, rewriter, loc);
       if (failed(inputWindow))
@@ -2782,58 +2951,22 @@ static FailureOr<Value> createConvOutputFromNchwRowStripFragments(Value rowStrip
           if (failed(patchRow))
             return failure();
 
-          Value paddedRow = *patchRow;
-          if (patchSize != paddedK)
-            paddedRow = createZeroPaddedTensor(
-              paddedRow, paddedPatchRowType, {0, 0}, {0, paddedK - patchSize}, rewriter, widthLoc);
-
-          Value zeroRow = createZeroTensorConstant(paddedRowType, rewriter);
-          auto kLoop = buildNormalizedScfFor(
-            rewriter,
-            widthLoc,
-            c0,
-            cNumKSlices,
-            c1,
-            ValueRange {zeroRow},
-            [&](OpBuilder&, Location reduceLoc, Value kSlice, ValueRange reduceIterArgs, SmallVectorImpl<Value>& reduceYielded) {
-              Value kOffset = arith::MulIOp::create(rewriter, reduceLoc, kSlice, cXbar);
-              SmallVector<OpFoldResult> aOffsets {rewriter.getIndexAttr(0), kOffset};
-              SmallVector<OpFoldResult> aSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(xbarDim)};
-              Value aTile = tensor::ExtractSliceOp::create(
-                rewriter, reduceLoc, paddedRowType, paddedRow, aOffsets, aSizes, getUnitStrides(rewriter, 2));
-              SmallVector<OpFoldResult> bOffsets {kOffset, rewriter.getIndexAttr(0)};
-              SmallVector<OpFoldResult> bSizes {rewriter.getIndexAttr(xbarDim), rewriter.getIndexAttr(xbarDim)};
-              Value bTile = extractStaticSliceOrIdentity(rewriter,
-                                                         reduceLoc,
-                                                         args.weights.front(),
-                                                         paddedWeightTileType,
-                                                         bOffsets,
-                                                         bSizes,
-                                                         getUnitStrides(rewriter, 2));
-              Value piece = spatial::SpatVMMOp::create(rewriter, reduceLoc, paddedRowType, bTile, aTile).getResult();
-              reduceYielded.push_back(
-                spatial::SpatVAddOp::create(rewriter, reduceLoc, paddedRowType, reduceIterArgs.front(), piece).getResult());
-              return success();
-            });
-          if (failed(kLoop))
+          FailureOr<Value> outputRow = createPaddedConvOutputRow(*patchRow,
+                                                                 state,
+                                                                 args.weights.front(),
+                                                                 state.hasBias ? args.inputs[1] : Value(),
+                                                                 paddedK,
+                                                                 numKSlices,
+                                                                 xbarDim,
+                                                                 rewriter,
+                                                                 widthLoc);
+          if (failed(outputRow))
             return failure();
 
-          Value rowResult = kLoop->results.front();
-          if (state.hasBias)
-            rowResult = spatial::SpatVAddOp::create(rewriter, widthLoc, paddedRowType, rowResult, args.inputs[1]).getResult();
-          Value outputRow = rowResult;
-          if (state.numChannelsOut != xbarDim) {
-            SmallVector<OpFoldResult> outputOffsets {rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
-            SmallVector<OpFoldResult> outputSizes {
-              rewriter.getIndexAttr(1), rewriter.getIndexAttr(state.numChannelsOut)};
-            outputRow = tensor::ExtractSliceOp::create(
-              rewriter, widthLoc, rowType, rowResult, outputOffsets, outputSizes, getUnitStrides(rewriter, 2));
-          }
-
           Value outputFragment = tensor::ExpandShapeOp::create(rewriter,
                                                                widthLoc,
                                                                outputPixelType,
-                                                               outputRow,
+                                                               *outputRow,
                                                                SmallVector<ReassociationIndices> {{0}, {1, 2, 3}});
           SmallVector<OpFoldResult> rowOffsets {
             rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), widthIndex};
@@ -2915,72 +3048,6 @@ static Value createFragmentReciprocalConstant(const DistributedTensorStep& step,
                              fragmentType);
 }
 
-[[maybe_unused]] static FailureOr<Value> createConvRowsForStrategy(const ConvLoweringState& state,
-                                                                   const ConvLoweringDecision& decision,
-                                                                   PatternRewriter& rewriter,
-                                                                   Location loc) {
-  auto wDenseAttr = getHostConstDenseElementsAttr(state.w);
-  PreparedConvInput preparedInput = standard::prepareInputForIm2Col(state, rewriter, loc);
-  Value biasMatrix;
-  DenseElementsAttr biasDenseAttr;
-  if (state.hasBias) {
-    biasDenseAttr = getHostConstDenseElementsAttr(state.b);
-    biasMatrix = expandBiasIfNeeded(state.b, rewriter, loc);
-  }
-
-  switch (decision.strategy) {
-  case PimConvLoweringLegacy:
-  case PimConvLoweringPackedIm2Col: {
-    standard::ConvGemmPlan plan = standard::buildConvGemmPlan(
-      state, static_cast<bool>(wDenseAttr), !state.hasBias || static_cast<bool>(biasDenseAttr), 0,
-      state.batchSize * state.outHeight * state.outWidth);
-    Value weightMatrix = standard::createWeightMatrix(state.w, plan, rewriter, loc);
-    Value gemmInputRows = standard::createIm2colRows(state, preparedInput, plan, rewriter, loc);
-    Value gemmB = standard::buildPackedWeights(wDenseAttr, weightMatrix, state, plan, rewriter, loc);
-    Value gemmBias = createZeroGemmBias(plan.gemmOutputRowsType, rewriter);
-    if (state.hasBias)
-      gemmBias = state.b;
-    Value gemmC = standard::buildPackedBias(gemmBias, biasMatrix, biasDenseAttr, state, plan, rewriter, loc);
-    Value gemmRows = ONNXGemmOp::create(rewriter,
-                                        loc,
-                                        plan.gemmOutputRowsType,
-                                        gemmInputRows,
-                                        gemmB,
-                                        gemmC,
-                                        APFloat(1.0f),
-                                        APFloat(1.0f),
-                                        /*transA=*/0,
-                                        /*transB=*/0)
-                       .getY();
-    return standard::maybeUnpackChunkRows(gemmRows, plan, rewriter, loc);
-  }
-  case PimConvLoweringStreamedPatch:
-  case PimConvLoweringOutputChannelTiled:
-  case PimConvLoweringTiled2D:
-  case PimConvLoweringStreamedPacked: {
-    standard::ConvGemmPlan seedPlan = standard::buildConvGemmPlan(
-      state, static_cast<bool>(wDenseAttr), !state.hasBias || static_cast<bool>(biasDenseAttr), 0, 1,
-      decision.strategy == PimConvLoweringStreamedPacked ? buildConvGeometry(state).pack : 1);
-    Value weightMatrix = standard::createWeightMatrix(state.w, seedPlan, rewriter, loc);
-    ConvGeometry geo = buildConvGeometry(state);
-    int64_t packFactor = decision.strategy == PimConvLoweringStreamedPacked ? geo.pack : 1;
-    uint64_t chunkPositions = chooseStreamChunkPositions(geo, packFactor);
-    return standard::createChunkedConvRows(state,
-                                           preparedInput,
-                                           weightMatrix,
-                                           biasMatrix,
-                                           wDenseAttr,
-                                           biasDenseAttr,
-                                           packFactor,
-                                           chunkPositions,
-                                           rewriter,
-                                           loc);
-  }
-  default:
-    return failure();
-  }
-}
-
 [[maybe_unused]] static FailureOr<DistributedTensorInfo> applyDistributedPreservingStep(const DistributedTensorInfo& inputInfo,
                                                                                         const DistributedTensorStep& step,
                                                                                         PatternRewriter& rewriter,
@@ -3741,6 +3808,8 @@ LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp) {
     return failure();
   if (state->outType.getRank() != 4 || !state->outType.hasStaticShape())
     return failure();
+  if (!getHostConstDenseElementsAttr(state->w))
+    return failure();
   if (state->hasBias && !isSupportedBiasAddValue(state->b, state->outType))
     return failure();
 
@@ -3752,6 +3821,8 @@ LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp) {
   analysis.barrierKind = DistributedConvBarrierKind::UnsupportedConsumer;
   analysis.barrierDetail = "selected row-strip layout";
   ConvGeometry geometry = buildConvGeometry(*state);
+  if (geometry.c > geometry.xbarSize)
+    return failure();
   ConvLoweringDecision decision = chooseConvLoweringStrategy(geometry, *requestedStrategy, analysis);
   if (decision.strategy == PimConvLoweringDepthwise && !depthwise::canUseStructuredRewrite(*state)
       && *requestedStrategy == PimConvLoweringAuto) {
@@ -3830,21 +3901,15 @@ lowerSelectedConv2DPlan(spatial::SpatConv2DPlanOp planOp,
     const bool applyBiasAfterStorage = rowState.hasBias;
     Value originalBias = rowState.b;
     if (applyBiasAfterStorage) {
-      if (!isSupportedBiasAddValue(originalBias, rowState.outType))
-        return planOp.emitOpError("selected row-strip Conv bias must be host-constant scalar/per-channel NCHW"),
-               failure();
       rowState.b = Value();
       rowState.hasBias = false;
     }
 
-    FailureOr<Value> rows = createConvRowsForStrategy(rowState, decision, rewriter, planOp.getLoc());
-    if (failed(rows))
-      return failure();
-    FailureOr<Value> rowStripStorage = createRowStripStorageFromRows(*rows, state->outType, rewriter, planOp.getLoc());
+    FailureOr<Value> rowStripStorage = createRowStripConvOutputFromDenseInput(rowState, rewriter, planOp.getLoc());
     if (failed(rowStripStorage))
       return planOp.emitOpError("failed to build row-strip fragment storage for the selected Conv plan"), failure();
     if (applyBiasAfterStorage) {
-      rowStripStorage = applyRowStripBiasAdd(*rowStripStorage, rowState.outType, originalBias, rewriter, planOp.getLoc());
+      rowStripStorage = applyRowStripBiasAdd(*rowStripStorage, state->outType, originalBias, rewriter, planOp.getLoc());
       if (failed(rowStripStorage))
         return planOp.emitOpError("failed to apply row-strip Conv bias per fragment"), failure();
     }

src/PIM/Dialect/Spatial/CMakeLists.txt

@@ -8,6 +8,7 @@ add_pim_library(SpatialOps
   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

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

@@ -0,0 +1,240 @@
+#include "CommunicationPlan.hpp"
+
+#include "mlir/IR/Diagnostics.h"
+
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/TypeSwitch.h"
+
+using namespace mlir;
+
+namespace onnx_mlir::spatial {
+namespace {
+
+struct GroupSignature {
+  CommunicationPhase phase = CommunicationPhase::Local;
+  ClassId sourceClass = 0;
+  ClassId targetClass = 0;
+  Type payloadType;
+  bool sourceIsBatch = false;
+  bool targetIsBatch = false;
+
+  bool operator==(const GroupSignature& other) const {
+    return phase == other.phase && sourceClass == other.sourceClass && targetClass == other.targetClass
+        && payloadType == other.payloadType && sourceIsBatch == other.sourceIsBatch
+        && targetIsBatch == other.targetIsBatch;
+  }
+};
+
+struct GroupSignatureInfo {
+  static GroupSignature getEmptyKey() {
+    return {CommunicationPhase::Local,
+            llvm::DenseMapInfo<ClassId>::getEmptyKey(),
+            llvm::DenseMapInfo<ClassId>::getEmptyKey(),
+            llvm::DenseMapInfo<Type>::getEmptyKey(),
+            false,
+            false};
+  }
+
+  static GroupSignature getTombstoneKey() {
+    return {CommunicationPhase::Local,
+            llvm::DenseMapInfo<ClassId>::getTombstoneKey(),
+            llvm::DenseMapInfo<ClassId>::getTombstoneKey(),
+            llvm::DenseMapInfo<Type>::getTombstoneKey(),
+            false,
+            false};
+  }
+
+  static unsigned getHashValue(const GroupSignature& key) {
+    return llvm::hash_combine(static_cast<unsigned>(key.phase),
+                              key.sourceClass,
+                              key.targetClass,
+                              key.payloadType,
+                              key.sourceIsBatch,
+                              key.targetIsBatch);
+  }
+
+  static bool isEqual(const GroupSignature& lhs, const GroupSignature& rhs) { return lhs == rhs; }
+};
+
+void emitPlanError(Operation* anchor, StringRef message) {
+  if (anchor)
+    anchor->emitError(message);
+}
+
+} // namespace
+
+FailureOr<CommunicationPlan>
+CommunicationPlan::buildProjectedInputPlan(Operation* anchor,
+                                           ArrayRef<CommunicationClassInfo> classes,
+                                           const DenseMap<CpuId, ClassId>& cpuToClass,
+                                           const ProjectedInputPlanMap& projectedPlans) {
+  CommunicationPlan plan;
+
+  DenseMap<ClassId, CommunicationClassInfo> classById;
+  for (CommunicationClassInfo klass : classes)
+    classById[klass.classId] = klass;
+
+  DenseSet<int64_t> remoteChannels;
+  DenseSet<const ProjectedInputTransferFragment*> seenFragments;
+  DenseMap<GroupSignature, unsigned, GroupSignatureInfo> groupBySignature;
+  unsigned nextPlaceholder = 0;
+
+  for (const auto& extractEntry : projectedPlans) {
+    for (const auto& classEntry : extractEntry.second) {
+      ClassId targetClass = classEntry.first;
+      const 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 (const 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;
+
+        PlaceholderId placeholder {nextPlaceholder++};
+        plan.demands.push_back(InputDemand {placeholder, fragment.producer, targetClass, inputPlan.layout.fragmentType});
+        plan.publications.push_back(
+          ProducerPublication {fragment.producer, sourceClass, inputPlan.layout.fragmentType});
+
+        if (fragment.sourceCoreId == fragment.targetCoreId) {
+          plan.phaseByFragment[&fragment] = CommunicationPhase::Local;
+          continue;
+        }
+
+        if (sourceInfo.rank == targetInfo.rank && sourceClass != targetClass) {
+          emitPlanError(anchor, "communication plan found a same-rank inter-class projected exchange");
+          return failure();
+        }
+
+        if (!remoteChannels.insert(fragment.channelId).second) {
+          emitPlanError(anchor, "communication plan found a duplicate remote projected channel id");
+          return failure();
+        }
+
+        TransferKind kind = sourceClass == targetClass ? TransferKind::SameClassForward : TransferKind::RemoteChannel;
+        CommunicationPhase phase =
+          sourceInfo.rank == targetInfo.rank
+            ? (fragment.sourceCoreId < fragment.targetCoreId ? CommunicationPhase::LowToHigh
+                                                             : CommunicationPhase::HighToLow)
+            : (sourceInfo.rank < targetInfo.rank ? CommunicationPhase::LowToHigh : CommunicationPhase::HighToLow);
+        plan.phaseByFragment[&fragment] = phase;
+
+        size_t exchangeId = plan.exchanges.size();
+        plan.exchangeByFragment[&fragment] = exchangeId;
+        plan.exchanges.push_back(ExchangeDescriptor {
+          exchangeId,
+          kind,
+          placeholder,
+          fragment.producer,
+          sourceClass,
+          targetClass,
+          sourceInfo.rank,
+          targetInfo.rank,
+          inputPlan.layout.fragmentType,
+          fragment.channelId,
+          fragment.sourceCoreId,
+          fragment.targetCoreId,
+          fragment.targetLane,
+          fragment.ordinal,
+          phase,
+          &fragment,
+        });
+
+        GroupSignature signature {
+          phase,
+          sourceClass,
+          targetClass,
+          inputPlan.layout.fragmentType,
+          sourceInfo.isBatch,
+          targetInfo.isBatch,
+        };
+        auto groupIt = groupBySignature.find(signature);
+        if (groupIt == groupBySignature.end()) {
+          unsigned groupIndex = static_cast<unsigned>(plan.batchGroups.size());
+          groupBySignature[signature] = groupIndex;
+          plan.batchGroups.push_back(CommunicationBatchGroup {
+            phase,
+            sourceClass,
+            targetClass,
+            inputPlan.layout.fragmentType,
+            sourceInfo.isBatch,
+            targetInfo.isBatch,
+            {},
+          });
+          groupIt = groupBySignature.find(signature);
+        }
+        plan.batchGroups[groupIt->second].exchanges.push_back(exchangeId);
+      }
+    }
+  }
+
+  for (const ExchangeDescriptor& exchange : plan.exchanges) {
+    if (!exchange.payloadType) {
+      emitPlanError(anchor, "communication plan has an exchange without payload type");
+      return failure();
+    }
+    if (exchange.phase == CommunicationPhase::Local) {
+      emitPlanError(anchor, "communication plan classified a remote exchange as local");
+      return failure();
+    }
+    if (exchange.kind == TransferKind::SameClassForward && exchange.sourceClass != exchange.targetClass) {
+      emitPlanError(anchor, "communication plan same-class exchange crosses class ids");
+      return failure();
+    }
+    if (exchange.kind == TransferKind::RemoteChannel && exchange.sourceRank == exchange.targetRank) {
+      emitPlanError(anchor, "communication plan remote exchange has equal source and target rank");
+      return failure();
+    }
+    if (exchange.sourceRank < exchange.targetRank && exchange.phase != CommunicationPhase::LowToHigh) {
+      emitPlanError(anchor, "communication plan phase does not match low-to-high rank direction");
+      return failure();
+    }
+    if (exchange.sourceRank > exchange.targetRank && exchange.phase != CommunicationPhase::HighToLow) {
+      emitPlanError(anchor, "communication plan phase does not match high-to-low rank direction");
+      return failure();
+    }
+    if (!exchange.fragment) {
+      emitPlanError(anchor, "communication plan exchange is not bound to a projected fragment");
+      return failure();
+    }
+  }
+
+  return plan;
+}
+
+std::optional<CommunicationPhase>
+CommunicationPlan::getPhase(const ProjectedInputTransferFragment& fragment) const {
+  auto it = phaseByFragment.find(&fragment);
+  if (it == phaseByFragment.end())
+    return std::nullopt;
+  return it->second;
+}
+
+std::optional<size_t> CommunicationPlan::getExchangeId(const ProjectedInputTransferFragment& fragment) const {
+  auto it = exchangeByFragment.find(&fragment);
+  if (it == exchangeByFragment.end())
+    return std::nullopt;
+  return it->second;
+}
+
+} // namespace onnx_mlir::spatial

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

@@ -0,0 +1,121 @@
+#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/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 CommunicationPhase {
+  Local,
+  LowToHigh,
+  HighToLow
+};
+
+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;
+};
+
+struct ExchangeDescriptor {
+  size_t id = 0;
+  TransferKind kind = TransferKind::RemoteChannel;
+  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 targetLane = 0;
+  unsigned ordinal = 0;
+  CommunicationPhase phase = CommunicationPhase::Local;
+  const ProjectedInputTransferFragment* fragment = nullptr;
+};
+
+struct CommunicationBatchGroup {
+  CommunicationPhase phase = CommunicationPhase::Local;
+  ClassId sourceClass = 0;
+  ClassId targetClass = 0;
+  mlir::Type payloadType;
+  bool sourceIsBatch = false;
+  bool targetIsBatch = false;
+  llvm::SmallVector<size_t, 8> exchanges;
+};
+
+class CommunicationPlan {
+public:
+  using ProjectedInputPlanMap =
+    llvm::DenseMap<mlir::Operation*, llvm::DenseMap<ClassId, ProjectedInputTransferPlan>>;
+
+  static mlir::FailureOr<CommunicationPlan>
+  buildProjectedInputPlan(mlir::Operation* anchor,
+                          llvm::ArrayRef<CommunicationClassInfo> classes,
+                          const llvm::DenseMap<CpuId, ClassId>& cpuToClass,
+                          const ProjectedInputPlanMap& projectedPlans);
+
+  std::optional<CommunicationPhase> getPhase(const ProjectedInputTransferFragment& fragment) const;
+  std::optional<size_t> getExchangeId(const ProjectedInputTransferFragment& fragment) const;
+  const ExchangeDescriptor& getExchange(size_t exchangeId) const { return exchanges[exchangeId]; }
+  llvm::ArrayRef<ExchangeDescriptor> getExchanges() const { return exchanges; }
+  llvm::ArrayRef<CommunicationBatchGroup> getBatchGroups() const { return batchGroups; }
+
+private:
+  llvm::DenseMap<const ProjectedInputTransferFragment*, CommunicationPhase> phaseByFragment;
+  llvm::DenseMap<const ProjectedInputTransferFragment*, size_t> exchangeByFragment;
+  llvm::SmallVector<InputDemand, 16> demands;
+  llvm::SmallVector<ProducerPublication, 16> publications;
+  llvm::SmallVector<ExchangeDescriptor, 16> exchanges;
+  llvm::SmallVector<CommunicationBatchGroup, 8> batchGroups;
+};
+
+} // namespace onnx_mlir::spatial

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

@@ -12,6 +12,7 @@
 #include <optional>
 
 #include "MaterializeMergeSchedule.hpp"
+#include "CommunicationPlan.hpp"
 #include "MergeMessages.hpp"
 #include "MergeScheduleKeys.hpp"
 #include "ProjectedFragments.hpp"
@@ -159,13 +160,6 @@ struct PendingProjectedScalarSend {
   mlir::Location loc;
 };
 
-struct PendingProjectedInputSend {
-  ClassId sourceClass = 0;
-  mlir::Value payload;
-  llvm::SmallVector<ProjectedInputTransferFragment*, 4> fragments;
-  mlir::Location loc;
-};
-
 enum class BatchInputDemandKind {
   LaneFragment,
   ProjectedFragment,
@@ -177,6 +171,14 @@ struct BatchInputDemand {
   std::optional<ProducerKey> wholeTensorProducer;
 };
 
+struct ProjectedExchangeEmissionState {
+  bool producerBound = false;
+  bool sendEmitted = false;
+  bool receiveEmitted = false;
+  mlir::Value producerPayload;
+  std::optional<mlir::Location> loc;
+};
+
 struct CloneIndexingContext {
   std::optional<mlir::Value> runSlotIndex;
   std::optional<mlir::Value> projectionSlotIndex;
@@ -267,6 +269,7 @@ struct MaterializerState {
   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>
@@ -282,6 +285,11 @@ struct MaterializerState {
     projectedExtractReplacements;
   llvm::DenseMap<mlir::Operation*, llvm::DenseMap<ClassId, ProjectedInputTransferPlan>>
     projectedInputTransferPlans;
+  CommunicationPlan projectedInputCommunicationPlan;
+  llvm::SmallVector<ProjectedExchangeEmissionState, 16> projectedExchangeStates;
+  CommunicationPhase currentProjectedCommunicationPhase = CommunicationPhase::LowToHigh;
+  llvm::SmallVector<size_t, 8> deferredProjectedExchanges;
+  llvm::SmallVector<ProducerKey, 8> deferredProducerValues;
   AvailableValueStore availableValues;
   llvm::DenseMap<mlir::Value, mlir::Value> hostReplacements;
   llvm::DenseMap<mlir::Value, ClassId> hostOutputOwners;
@@ -290,9 +298,6 @@ struct MaterializerState {
   llvm::SmallVector<ScalarPeerEdgeKey, 8> materializedScalarPeerReceives;
   llvm::SmallVector<PendingScalarSend, 8> pendingScalarSends;
   llvm::SmallVector<PendingProjectedScalarSend, 8> pendingProjectedScalarSends;
-  llvm::SmallVector<PendingProjectedInputSend, 16> pendingProjectedInputSends;
-  llvm::DenseSet<ClassId> projectedInputPhaseBarrierClasses;
-  llvm::DenseMap<ClassId, unsigned> pendingProjectedHighToLowReceives;
   llvm::DenseSet<mlir::Operation*> oldComputeOps;
 
   MaterializerState(mlir::func::FuncOp func, const MergeScheduleResult& schedule, int64_t& nextChannelId)