meno diamantini

src/PIM/Dialect/Spatial/Spatial.td

@@ -232,6 +232,22 @@ def SpatReluPlanOp : SpatOp<"relu_plan", []> {
   let hasVerifier = 1;
 }
 
+def SpatBiasAddPlanOp : SpatOp<"bias_add_plan", []> {
+  let summary = "Layout-aware Conv-style bias add planning op";
+
+  let arguments = (ins
+    SpatTensor:$input,
+    SpatTensor:$bias,
+    StrAttr:$logicalLayout
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let hasVerifier = 1;
+}
+
 def SpatBlueprintOp : SpatOp<"blueprint", []> {
   let summary = "Blueprint for assembling logical tensors from published fragments";
 

src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp

@@ -1,5 +1,6 @@
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/Block.h"
@@ -59,6 +60,21 @@ static LogicalResult verifyStaticWeights(ComputeOpTy computeOp, StringRef kind)
   return success();
 }
 
+static bool isStaticScfForInductionVar(Value value) {
+  auto blockArg = dyn_cast<BlockArgument>(value);
+  if (!blockArg)
+    return false;
+
+  auto loop = dyn_cast_or_null<scf::ForOp>(blockArg.getOwner()->getParentOp());
+  if (!loop || loop.getInductionVar() != value)
+    return false;
+
+  std::optional<int64_t> lowerBound = matchConstantIndexValue(loop.getLowerBound());
+  std::optional<int64_t> upperBound = matchConstantIndexValue(loop.getUpperBound());
+  std::optional<int64_t> step = matchConstantIndexValue(loop.getStep());
+  return lowerBound && upperBound && step && *step > 0 && *upperBound >= *lowerBound;
+}
+
 static bool isStaticIndexExpr(Value value) {
   if (matchConstantIndexValue(value))
     return true;
@@ -80,7 +96,7 @@ static bool isStaticIndexExpr(Value value) {
 }
 
 static bool isSupportedLaneOffsetExpr(Value value, BlockArgument laneArg) {
-  if (value == laneArg || isStaticIndexExpr(value))
+  if (value == laneArg || isStaticIndexExpr(value) || isStaticScfForInductionVar(value))
     return true;
 
   auto affineApply = value.getDefiningOp<affine::AffineApplyOp>();
@@ -436,6 +452,39 @@ LogicalResult SpatReluPlanOp::verify() {
   return success();
 }
 
+LogicalResult SpatBiasAddPlanOp::verify() {
+  if (failed(verifyPlanTensorTypes(getOperation(), getInput(), getOutput(), "spat.bias_add_plan")))
+    return failure();
+  if (!isKnownLogicalLayout(getLogicalLayout()))
+    return emitError("requires a known logical layout");
+
+  auto inputType = dyn_cast<RankedTensorType>(getInput().getType());
+  auto biasType = dyn_cast<RankedTensorType>(getBias().getType());
+  auto outputType = dyn_cast<RankedTensorType>(getOutput().getType());
+  if (!inputType || !biasType || !outputType)
+    return emitError("requires ranked tensor input, bias, and output");
+  if (!inputType.hasStaticShape() || !biasType.hasStaticShape() || !outputType.hasStaticShape())
+    return emitError("requires static tensor input, bias, and output");
+  if (inputType != outputType)
+    return emitError("requires matching input and output tensor types");
+  if (outputType.getRank() != 4)
+    return emitError("requires rank-4 input/output tensors");
+  if (getLogicalLayout() != "nchw")
+    return emitError("requires logical layout \"nchw\"");
+  if (biasType.getElementType() != outputType.getElementType())
+    return emitError("requires bias element type to match the output element type");
+
+  ArrayRef<int64_t> biasShape = biasType.getShape();
+  const int64_t channels = outputType.getDimSize(1);
+  const bool supported = biasShape.empty() || (biasShape.size() == 1 && biasShape[0] == channels)
+                         || (biasShape.size() == 2 && biasShape[0] == 1 && biasShape[1] == channels)
+                         || (biasShape.size() == 4 && biasShape[0] == 1 && biasShape[1] == channels
+                             && biasShape[2] == 1 && biasShape[3] == 1);
+  if (!supported)
+    return emitError("requires scalar or per-channel bias broadcastable over NCHW");
+  return success();
+}
+
 LogicalResult SpatBlueprintOp::verify() {
   auto modeAttr = getModeAttr();
   bool isFragmentAssembly = modeAttr && modeAttr.getValue() == "fragment_assembly";

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;