A new Beginning

2026-07-07 18:28:37 +02:00
parent 8338caf3f3
commit 1f4f58de1c
25 changed files with 1187 additions and 13533 deletions
@@ -50,6 +50,64 @@ static void printBlockArgumentList(OpAsmPrinter& printer, ArrayRef<BlockArgument
  printer << ")";
 }

+static void printBlockHeaderWithoutArgLocs(OpAsmPrinter& printer, Block& block) {
+  printer.printSuccessor(&block);
+  if (block.getNumArguments() == 0) {
+    printer << ":";
+    return;
+  }
+
+  printer << "(";
+  for (auto [index, argument] : llvm::enumerate(block.getArguments())) {
+    if (index != 0)
+      printer << ", ";
+    printer.printOperand(argument);
+    printer << ": ";
+    printer.printType(argument.getType());
+  }
+  printer << "):";
+}
+
+static void printRegionWithoutBlockArgLocs(OpAsmPrinter& printer,
+                                           Region& region,
+                                           bool printEntryBlockArgs = true,
+                                           bool printBlockTerminators = true,
+                                           bool printEmptyBlock = false,
+                                           bool printEntryBlockHeaderWhenMultiblock = false) {
+  printer << " {";
+  if (region.empty()) {
+    if (printEmptyBlock)
+      printer << "\n";
+    printer << "}";
+    return;
+  }
+
+  bool hasMultipleBlocks = std::next(region.begin()) != region.end();
+  for (auto [blockIndex, block] : llvm::enumerate(region)) {
+    bool printBlockHeader = blockIndex != 0 || printEntryBlockArgs
+                            || (blockIndex == 0 && printEntryBlockHeaderWhenMultiblock && hasMultipleBlocks)
+                            || (printEmptyBlock && block.empty());
+    unsigned indent = printBlockHeader ? 4u : 2u;
+
+    if (printBlockHeader) {
+      printer.getStream() << "\n";
+      printer.getStream().indent(2);
+      printBlockHeaderWithoutArgLocs(printer, block);
+    }
+
+    for (Operation& nestedOp : block) {
+      if (!printBlockTerminators && nestedOp.hasTrait<OpTrait::IsTerminator>())
+        continue;
+      printer.getStream() << "\n";
+      printer.getStream().indent(indent);
+      printer.printCustomOrGenericOp(&nestedOp);
+    }
+  }
+
+  printer.getStream() << "\n";
+  printer << "}";
+}
+
 static ParseResult parseBlockArgumentList(OpAsmParser& parser, SmallVectorImpl<OpAsmParser::Argument>& arguments) {
  if (parser.parseLParen())
    return failure();
@@ -160,7 +218,9 @@ void printComputeLikeOp(ComputeOpTy op, OpAsmPrinter& printer) {
  printer << " -> ";
  printCompressedTypeSequence(printer, op.getResultTypes());
  printer << " ";
-  printer.printRegion(op.getBody(), /*printEntryBlockArgs=*/false);
+  printRegionWithoutBlockArgLocs(
+    printer, op.getBody(), /*printEntryBlockArgs=*/false, /*printBlockTerminators=*/true,
+    /*printEmptyBlock=*/false, /*printEntryBlockHeaderWhenMultiblock=*/true);
 }

 template <typename ComputeOpTy>
@@ -290,7 +350,9 @@ void printComputeBatchLikeOp(ComputeBatchOpTy op, OpAsmPrinter& printer) {
  printer << " -> ";
  printCompressedTypeSequence(printer, op.getResultTypes());
  printer << " ";
-  printer.printRegion(op.getBody(), /*printEntryBlockArgs=*/false);
+  printRegionWithoutBlockArgLocs(
+    printer, op.getBody(), /*printEntryBlockArgs=*/false, /*printBlockTerminators=*/true,
+    /*printEmptyBlock=*/false, /*printEntryBlockHeaderWhenMultiblock=*/true);
 }

 template <typename ComputeBatchOpTy>
@@ -407,6 +469,89 @@ ParseResult SpatYieldOp::parse(OpAsmParser& parser, OperationState& result) {
  return parser.resolveOperands(outputs, outputTypes, parser.getCurrentLocation(), result.operands);
 }

+void SpatBlockYieldOp::print(OpAsmPrinter& printer) {
+  printer << " ";
+  printCompressedValueSequence(printer, getOutputs());
+  if (getOperation()->getNumSuccessors() != 0) {
+    printer << " next ";
+    printer.printSuccessor(getOperation()->getSuccessor(0));
+  }
+  printer.printOptionalAttrDict((*this)->getAttrs());
+  printer << " : ";
+  printCompressedTypeSequence(printer, getOutputs().getTypes());
+}
+
+ParseResult SpatBlockYieldOp::parse(OpAsmParser& parser, OperationState& result) {
+  SmallVector<OpAsmParser::UnresolvedOperand> outputs;
+  SmallVector<Type> outputTypes;
+  Block* successor = nullptr;
+
+  OpAsmParser::UnresolvedOperand firstOutput;
+  OptionalParseResult firstOutputResult = parser.parseOptionalOperand(firstOutput);
+  if (firstOutputResult.has_value()) {
+    if (failed(*firstOutputResult))
+      return failure();
+    if (parseCompressedOperandEntryWithFirst(parser, firstOutput, outputs))
+      return failure();
+    while (succeeded(parser.parseOptionalComma()))
+      if (parseOneCompressedOperandEntry(parser, outputs))
+        return failure();
+  }
+
+  if (succeeded(parser.parseOptionalKeyword("next")) && parser.parseSuccessor(successor))
+    return failure();
+
+  if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()
+      || parseCompressedTypeSequence(parser, outputTypes, /*allowEmpty=*/true))
+    return failure();
+
+  if (outputs.size() != outputTypes.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of outputs and output types must match");
+  if (parser.resolveOperands(outputs, outputTypes, parser.getCurrentLocation(), result.operands))
+    return failure();
+  if (successor)
+    result.addSuccessors(successor);
+  return success();
+}
+
+void SpatDeferredCommunicationOp::print(OpAsmPrinter& printer) {
+  printer << " ";
+  printCompressedValueSequence(printer, getSources());
+  printer.printOptionalAttrDict((*this)->getAttrs());
+  printer << " : ";
+  printer.printFunctionalType(getSources().getTypes(), getOperation()->getResultTypes());
+  printer << " ";
+  printRegionWithoutBlockArgLocs(printer, getBody(), /*printEntryBlockArgs=*/false);
+}
+
+ParseResult SpatDeferredCommunicationOp::parse(OpAsmParser& parser, OperationState& result) {
+  SmallVector<OpAsmParser::UnresolvedOperand> sources;
+  Type functionTypeStorage;
+
+  if (parseCompressedOperandSequence(parser, sources) || parser.parseOptionalAttrDict(result.attributes)
+      || parser.parseColon() || parser.parseType(functionTypeStorage))
+    return failure();
+
+  auto functionType = dyn_cast<FunctionType>(functionTypeStorage);
+  if (!functionType)
+    return parser.emitError(parser.getCurrentLocation(), "expected deferred communication function type");
+  if (sources.size() != functionType.getNumInputs())
+    return parser.emitError(parser.getCurrentLocation(), "number of sources and source types must match");
+
+  if (parser.resolveOperands(sources, functionType.getInputs(), parser.getCurrentLocation(), result.operands))
+    return failure();
+  result.addTypes(functionType.getResults());
+
+  Region* body = result.addRegion();
+  SmallVector<OpAsmParser::Argument> bodyArgs;
+  for (Type type : functionType.getInputs()) {
+    OpAsmParser::Argument argument;
+    argument.type = type;
+    bodyArgs.push_back(argument);
+  }
+  return parser.parseRegion(*body, bodyArgs);
+}
+
 void SpatExtractRowsOp::print(OpAsmPrinter& printer) {
  printer << " ";
  printer.printOperand(getInput());
@@ -618,7 +763,8 @@ ParseResult SpatScheduledComputeBatch::parse(OpAsmParser& parser, OperationState

 void SpatInParallelOp::print(OpAsmPrinter& printer) {
  printer << " ";
-  printer.printRegion(getRegion(), /*printEntryBlockArgs=*/false, /*printBlockTerminators=*/false);
+  printRegionWithoutBlockArgLocs(
+    printer, getRegion(), /*printEntryBlockArgs=*/false, /*printBlockTerminators=*/false);
  printer.printOptionalAttrDict((*this)->getAttrs());
 }