A new Beginning

2026-07-07 18:28:37 +02:00
parent 8338caf3f3
commit 1f4f58de1c
25 changed files with 1187 additions and 13533 deletions
@@ -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 (auto extractSlice = dyn_cast<tensor::ExtractSliceOp>(&bodyOp))
-      if (failed(verifyStaticUnitStrideExtractSliceOp(extractSlice, *laneArg, "tensor.extract_slice")))
-        return failure();
-  }
+  if (verifyLaneSliceOffsets)
+    for (auto& bodyOp : block) {
+      if (auto extractSlice = dyn_cast<tensor::ExtractSliceOp>(&bodyOp))
+        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)
-    return compute.emitOpError("compute body must have weight and input block arguments");
+  bool isScheduled = isa<SpatScheduledCompute>(compute.getOperation());
+  if (compute.getBody().empty())
+    return compute.emitOpError("compute body must have at least one block");

-  for (auto [weightIndex, weight] : llvm::enumerate(compute.getWeights())) {
-    auto blockArg = compute.getWeightArgument(weightIndex);
-    if (!blockArg || blockArg->getType() != weight.getType())
-      return compute.emitOpError("compute weight block argument types must match weight operand types exactly");
-  }
-  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");
-  }
+  SmallVector<Type> yieldedTypes;
+  for (Block& block : compute.getBody()) {
+    if ((!isScheduled && block.getNumArguments() != expectedArgCount)
+        || (isScheduled && block.getNumArguments() < expectedArgCount))
+      return compute.emitOpError("compute body must have weight and input block arguments");

-  if (block.mightHaveTerminator()) {
-    auto yieldOp = dyn_cast_or_null<SpatYieldOp>(block.getTerminator());
-    if (!yieldOp)
+    for (auto [weightIndex, weight] : llvm::enumerate(compute.getWeights()))
+      if (block.getArgument(weightIndex).getType() != weight.getType())
+        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 yieldTypes = yieldOp->getOperandTypes();
-    if (resultTypes.size() != yieldTypes.size())
-      return compute.emitOpError("ComputeOp must have same number of results as yieldOp operands");
+  auto resultTypes = compute.getResultTypes();
+  if (resultTypes.size() != yieldedTypes.size())
+    return compute.emitOpError("ComputeOp must have same number of results as yielded operands");

-    for (auto it : llvm::reverse(llvm::zip(resultTypes, yieldTypes))) {
-      auto resultType = std::get<0>(it);
-      auto yieldType = std::get<1>(it);
+  for (auto it : llvm::reverse(llvm::zip(resultTypes, yieldedTypes))) {
+    auto resultType = std::get<0>(it);
+    auto yieldType = std::get<1>(it);

-      if (resultType != yieldType || failed(verifyCompatibleShape(resultType, yieldType)))
-        return compute.emitOpError("ComputeOp output must be of the same type as yieldOp operand");
+    if (resultType != yieldType || failed(verifyCompatibleShape(resultType, yieldType)))
+      return compute.emitOpError("ComputeOp output must be of the same type as yieldOp operand");

-      if (auto resultRankedType = dyn_cast<RankedTensorType>(resultType)) {
-        if (auto yieldRankedType = dyn_cast<RankedTensorType>(yieldType)) {
-          if (resultRankedType.getEncoding() != yieldRankedType.getEncoding())
-            return compute.emitOpError("ComputeOp output must have the same encoding as yieldOp operand");
-        }
-        else {
+    if (auto resultRankedType = dyn_cast<RankedTensorType>(resultType)) {
+      if (auto yieldRankedType = dyn_cast<RankedTensorType>(yieldType)) {
+        if (resultRankedType.getEncoding() != yieldRankedType.getEncoding())
          return compute.emitOpError("ComputeOp output has an encoding while yieldOp operand does not have one");
-        }
      }
-      else if (dyn_cast<RankedTensorType>(yieldType)) {
-        return compute.emitOpError("ComputeOp output must not have an encoding if yieldOp operand has one");
+      else {
+        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 (auto inputArg = compute.getInputArgument(inputIndex); !inputArg || inputArg->use_empty())
-      return compute.emitOpError("ComputeOp block argument is not used");
+  if (compute.getBody().hasOneBlock())
+    for (unsigned inputIndex = 0; inputIndex < compute.getInputs().size(); ++inputIndex)
+      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"); }