meno diamantini

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

@@ -121,6 +121,51 @@ struct CompactRunPlan {
   llvm::SmallVector<RunOutputDemand, 4> outputs;
 };
 
+struct ScalarPeerEdgeKey {
+  int64_t sourceCore = 0;
+  int64_t targetCore = 0;
+  mlir::Type payloadType;
+};
+
+struct ScalarPeerReceiveKey {
+  int64_t sourceCore = 0;
+  int64_t targetCore = 0;
+  mlir::Type payloadType;
+  std::optional<int64_t> channelId;
+};
+
+struct PendingScalarSend {
+  ClassId sourceClass = 0;
+  int64_t sourceCore = 0;
+  int64_t targetCore = 0;
+  mlir::Type payloadType;
+  ScalarPeerEdgeKey waitForReceive;
+  mlir::Operation* payloadAnchor = nullptr;
+  mlir::Value payload;
+  MessageVector messages;
+  mlir::Location loc;
+};
+
+struct PendingProjectedScalarSend {
+  ClassId sourceClass = 0;
+  int64_t sourceCore = 0;
+  int64_t targetCore = 0;
+  mlir::Type payloadType;
+  ScalarPeerEdgeKey waitForReceive;
+  mlir::Operation* payloadAnchor = nullptr;
+  mlir::Value payload;
+  MessageVector messages;
+  ProjectedTransferDescriptor descriptor;
+  mlir::Location loc;
+};
+
+struct PendingProjectedInputSend {
+  ClassId sourceClass = 0;
+  mlir::Value payload;
+  llvm::SmallVector<ProjectedInputTransferFragment*, 4> fragments;
+  mlir::Location loc;
+};
+
 enum class BatchInputDemandKind {
   LaneFragment,
   ProjectedFragment,
@@ -235,10 +280,19 @@ struct MaterializerState {
     projectedTransfers;
   llvm::DenseMap<mlir::Operation*, llvm::DenseMap<ClassId, ProjectedExtractReplacement>>
     projectedExtractReplacements;
+  llvm::DenseMap<mlir::Operation*, llvm::DenseMap<ClassId, ProjectedInputTransferPlan>>
+    projectedInputTransferPlans;
   AvailableValueStore availableValues;
   llvm::DenseMap<mlir::Value, mlir::Value> hostReplacements;
   llvm::DenseMap<mlir::Value, ClassId> hostOutputOwners;
   llvm::SmallVector<PendingProjectedHostOutputFragment, 32> pendingProjectedHostOutputFragments;
+  llvm::SmallVector<ScalarPeerEdgeKey, 16> plannedScalarPeerReceives;
+  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)

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

@@ -47,6 +47,24 @@ struct ProjectedExtractReplacement {
   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;