Unexpected invariant now it's clear (batched in the first tensor rank)

src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp

@@ -574,20 +574,26 @@ LogicalResult SpatBlueprintOp::verify() {
 
   auto stridesAttr = getFragmentStridesAttr();
   auto operandIndicesAttr = getFragmentOperandIndicesAttr();
+  auto sourceSlotsAttr = getFragmentSourceSlotsAttr();
   auto sourceOffsetsAttr = getFragmentSourceOffsetsAttr();
   if (!operandIndicesAttr)
     return emitError("fragment assembly blueprint requires fragment operand indices");
+  if (!sourceSlotsAttr)
+    return emitError("fragment assembly blueprint requires physical fragment source slots");
   if (!sourceOffsetsAttr)
     return emitError("fragment assembly blueprint requires fragment source offsets");
   if (!stridesAttr)
     return emitError("fragment assembly blueprint requires fragment strides");
   ArrayRef<int64_t> operandIndices = operandIndicesAttr.asArrayRef();
+  ArrayRef<int64_t> sourceSlots = sourceSlotsAttr.asArrayRef();
   ArrayRef<int64_t> sourceOffsets = sourceOffsetsAttr.asArrayRef();
   ArrayRef<int64_t> strides = stridesAttr.asArrayRef();
   if (strides.size() != offsets.size())
     return emitError("fragment stride and offset arrays must have the same length");
   if (sourceOffsets.size() != operandIndices.size())
     return emitError("fragment source offset count must match fragment operand index count");
+  if (sourceSlots.size() != operandIndices.size())
+    return emitError("fragment source slot count must match fragment operand index count");
   if (!getConflictPolicyAttr() || !getCoveragePolicyAttr())
     return emitError("fragment assembly blueprint requires conflict and coverage policies");
   if (getConflictPolicy() != "disjoint")
@@ -622,14 +628,19 @@ LogicalResult SpatBlueprintOp::verify() {
     int64_t operandIndex = operandIndices[fragmentIndex];
     if (operandIndex < 0 || operandIndex >= operandCount)
       return emitError("fragment assembly operand index is out of range");
+    if (sourceSlots[fragmentIndex] < 0)
+      return emitError("fragment assembly physical source slot must be nonnegative");
     if (sourceOffsets[fragmentIndex] < 0)
       return emitError("fragment assembly source offsets must be nonnegative");
 
     auto operandType = dyn_cast<RankedTensorType>(operands[operandIndex].getType());
     if (!operandType || !operandType.hasStaticShape())
       return emitError("fragment assembly blueprint requires static ranked tensor operands");
-    if (operandType.getRank() != rank)
-      return emitError("fragment assembly blueprint requires operand/result rank match");
+    if (operandType.getRank() != rank + 1)
+      return emitError("fragment assembly physical operand must have one leading source-slot dimension");
+    if (sourceSlots[fragmentIndex] >= operandType.getDimSize(0))
+      return emitError("fragment assembly physical source slot is out of range");
+    auto fragmentType = RankedTensorType::get(operandType.getShape().drop_front(), operandType.getElementType(), operandType.getEncoding());
 
     SmallVector<int64_t, 4> fragmentOffsets;
     SmallVector<int64_t, 4> fragmentSizes;
@@ -645,12 +656,12 @@ LogicalResult SpatBlueprintOp::verify() {
     int64_t fragmentElements = 1;
     for (int64_t dim = 0; dim < rank; ++dim)
       fragmentElements *= fragmentSizes[dim];
-    if (sourceOffsets[fragmentIndex] + fragmentElements > operandType.getNumElements())
-      return emitError("fragment assembly source offset exceeds the operand bounds");
+    if (sourceOffsets[fragmentIndex] + fragmentElements > fragmentType.getNumElements())
+      return emitError("fragment assembly source offset exceeds the selected physical fragment bounds");
     SmallVector<int64_t, 4> sourceSliceOffsets =
-      expandFlatElementIndex(sourceOffsets[fragmentIndex], operandType.getShape());
+      expandFlatElementIndex(sourceOffsets[fragmentIndex], fragmentType.getShape());
     for (int64_t dim = 0; dim < rank; ++dim)
-      if (sourceSliceOffsets[dim] + fragmentSizes[dim] > operandType.getDimSize(dim))
+      if (sourceSliceOffsets[dim] + fragmentSizes[dim] > fragmentType.getDimSize(dim))
         return emitError("fragment assembly source offset must describe a valid unit-stride slice");
 
     for (const auto& [existingOffsets, existingSizes] : slices) {
@@ -746,7 +757,8 @@ LogicalResult verifyComputeLikeOp(ComputeOpTy compute, StringRef opName) {
 
     Operation* terminator = block.getTerminator();
     if (auto yieldOp = dyn_cast_or_null<SpatYieldOp>(terminator)) {
-      if (isScheduled)
+      auto realized = compute->template getAttrOfType<BoolAttr>("scheduled.realized");
+      if (isScheduled && (!realized || !realized.getValue() || !compute.getBody().hasOneBlock()))
         return compute.emitOpError("scheduled compute blocks must terminate with spat.block_yield");
       llvm::append_range(yieldedTypes, yieldOp->getOperandTypes());
       continue;
@@ -820,6 +832,16 @@ LogicalResult SpatBlockYieldOp::verify() {
 LogicalResult SpatDeferredCommunicationOp::verify() {
   if (getSources().empty())
     return emitOpError("requires at least one source");
+  static constexpr StringLiteral staleAttributes[] = {
+    "exchangeId", "logicalProducer", "logicalConsumer", "sourceClass", "targetClass", "sourceCore",
+    "targetCore", "sourceLane", "targetLane", "transferKind", "resultIndex", "projectedTransfer",
+    "hostOutputOwner", "source_cpus", "source_classes", "source_lane_ranges", "target_cpus",
+    "target_classes", "target_lane_ranges", "batched", "source_operand_for_scheduled_lane",
+    "multi_source_payload"};
+  for (StringLiteral name : staleAttributes)
+    if (getOperation()->hasAttr(name))
+      return emitOpError() << "does not accept stale routing attribute '" << name
+                           << "'; source selection and shaping belong in the body and routing is derived in Phase 2";
   if (failed(verifyRegionArguments(getOperation(), getBody(), getSources(), "spat.deferred_communication")))
     return failure();
   return verifyYieldTypes(getOperation(), getBody(), getOperation()->getResultTypes(), "spat.deferred_communication");