refactorone

2026-05-20 19:06:41 +02:00
parent f56c4159b5
commit a50e77ff38
50 changed files with 3420 additions and 1187 deletions
@@ -20,6 +20,80 @@ static DenseI32ArrayAttr getDenseI32ArrayAttr(OpAsmParser& parser, ArrayRef<int3
  return parser.getBuilder().getDenseI32ArrayAttr(values);
 }

+static IntegerAttr getI32Attr(OpAsmParser& parser, int32_t value) {
+  return parser.getBuilder().getI32IntegerAttr(value);
+}
+
+static bool parseOptionalKeywordAlias(OpAsmParser& parser, StringRef preferred, StringRef legacy) {
+  return succeeded(parser.parseOptionalKeyword(preferred)) || succeeded(parser.parseOptionalKeyword(legacy));
+}
+
+static void printBlockArgumentList(OpAsmPrinter& printer, ArrayRef<BlockArgument> arguments) {
+  printer << "(";
+  for (auto [index, argument] : llvm::enumerate(arguments)) {
+    if (index != 0)
+      printer << ", ";
+    printer.printOperand(argument);
+  }
+  printer << ")";
+}
+
+static ParseResult parseBlockArgumentList(OpAsmParser& parser, SmallVectorImpl<OpAsmParser::Argument>& arguments) {
+  if (parser.parseLParen())
+    return failure();
+  if (succeeded(parser.parseOptionalRParen()))
+    return success();
+
+  OpAsmParser::Argument argument;
+  if (parser.parseArgument(argument))
+    return failure();
+  arguments.push_back(argument);
+  while (succeeded(parser.parseOptionalComma())) {
+    if (parser.parseArgument(argument))
+      return failure();
+    arguments.push_back(argument);
+  }
+  return parser.parseRParen();
+}
+
+static void printBoundValueList(OpAsmPrinter& printer, ValueRange arguments, ValueRange operands, ListDelimiter delimiter) {
+  printCompressedValueList(printer, arguments, delimiter);
+  printer << " = ";
+  printCompressedValueList(printer, operands, delimiter);
+}
+
+static ParseResult parseBoundValueList(OpAsmParser& parser,
+                                       ListDelimiter delimiter,
+                                       SmallVectorImpl<OpAsmParser::Argument>& arguments,
+                                       SmallVectorImpl<OpAsmParser::UnresolvedOperand>& operands) {
+  if (parseOpenDelimiter(parser, delimiter))
+    return failure();
+  if (succeeded(parseOptionalCloseDelimiter(parser, delimiter))) {
+    if (parser.parseEqual() || parseCompressedOperandList(parser, delimiter, operands))
+      return failure();
+    return success();
+  }
+
+  if (parseOneCompressedArgumentEntry(parser, arguments))
+    return failure();
+  while (succeeded(parser.parseOptionalComma()))
+    if (parseOneCompressedArgumentEntry(parser, arguments))
+      return failure();
+
+  auto parseCloseDelimiter = [&](ListDelimiter currentDelimiter) -> ParseResult {
+    switch (currentDelimiter) {
+    case ListDelimiter::Paren:
+      return parser.parseRParen();
+    case ListDelimiter::Square:
+      return parser.parseRSquare();
+    }
+    llvm_unreachable("unsupported delimiter");
+  };
+  if (parseCloseDelimiter(delimiter) || parser.parseEqual() || parseCompressedOperandList(parser, delimiter, operands))
+    return failure();
+  return success();
+}
+
 static void printCoreIdList(OpAsmPrinter& printer, StringRef keyword, ArrayRef<int32_t> coreIds) {
  printer << " " << keyword << " ";
  printCompressedIntegerList(printer, coreIds);
@@ -33,15 +107,76 @@ static ParseResult parseOptionalCoreIdList(OpAsmParser& parser, StringRef keywor

 } // namespace

-void PimCoreBatchOp::print(OpAsmPrinter& printer) {
-  printer << " lanes " << getLaneCount() << " ";
-  size_t weightsPerLane = getLaneCount() > 0 ? getWeights().size() / static_cast<size_t>(getLaneCount()) : 0;
-  if (getLaneCount() > 1 && hasRepeatedTuple(getWeights(), weightsPerLane))
-    printValueTupleRun(printer, getWeights(), weightsPerLane, ListDelimiter::Paren);
-  else
-    printCompressedValueList(printer, getWeights(), ListDelimiter::Paren);
+void PimCoreOp::print(OpAsmPrinter& printer) {
+  SmallVector<Value> weightArgs;
+  weightArgs.reserve(getWeights().size());
+  for (unsigned index = 0; index < getWeights().size(); ++index)
+    weightArgs.push_back(getWeightArgument(index));
+
  printer << " ";
-  printCompressedValueList(printer, getInputs(), ListDelimiter::Square);
+  printBoundValueList(printer, weightArgs, getWeights(), ListDelimiter::Square);
+  printer << " coreId " << getCoreId();
+  printer.printOptionalAttrDict((*this)->getAttrs(), {getCoreIdAttrName().getValue()});
+  printer << " : ";
+  printCompressedTypeList(printer, TypeRange(getWeights()), ListDelimiter::Square);
+  printer << " -> () ";
+  printer.printRegion(getBody(), /*printEntryBlockArgs=*/false);
+}
+
+ParseResult PimCoreOp::parse(OpAsmParser& parser, OperationState& result) {
+  SmallVector<OpAsmParser::Argument> weightArgs;
+  SmallVector<OpAsmParser::UnresolvedOperand> weights;
+  SmallVector<Type> weightTypes;
+  int32_t coreId = 0;
+
+  if (parseBoundValueList(parser, ListDelimiter::Square, weightArgs, weights))
+    return failure();
+
+  bool hasCoreId = parseOptionalKeywordAlias(parser, "coreId", "core_id");
+  if (hasCoreId && parser.parseInteger(coreId))
+    return failure();
+
+  if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()
+      || parseCompressedRepeatedList(
+        parser, ListDelimiter::Square, weightTypes, [&](Type& type) { return parser.parseType(type); })
+      || parser.parseArrow() || parser.parseLParen() || parser.parseRParen())
+    return failure();
+
+  if (weights.size() != weightTypes.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of weights and weight types must match");
+  if (weightArgs.size() != weights.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of weight bindings and weight operands must match");
+  if (hasCoreId && result.attributes.get("coreId"))
+    return parser.emitError(parser.getCurrentLocation(),
+                            "coreId cannot be specified both positionally and in attr-dict");
+
+  if (parser.resolveOperands(weights, weightTypes, parser.getCurrentLocation(), result.operands))
+    return failure();
+
+  if (hasCoreId)
+    result.addAttribute("coreId", getI32Attr(parser, coreId));
+
+  Region* body = result.addRegion();
+  applyArgumentTypes(weightTypes, weightArgs);
+  return parser.parseRegion(*body, weightArgs);
+}
+
+void PimCoreBatchOp::print(OpAsmPrinter& printer) {
+  printer << " ";
+  printer.printOperand(getLaneArgument());
+  printer << " = 0 to " << getLaneCount() << " ";
+
+  SmallVector<Value> weightArgs;
+  weightArgs.reserve(getWeights().size());
+  for (unsigned index = 0; index < getWeights().size(); ++index)
+    weightArgs.push_back(getWeightArgument(index));
+  printBoundValueList(printer, weightArgs, getWeights(), ListDelimiter::Square);
+  printer << " ";
+  SmallVector<Value> inputArgs;
+  inputArgs.reserve(getInputs().size());
+  for (unsigned index = 0; index < getInputs().size(); ++index)
+    inputArgs.push_back(getInputArgument(index));
+  printBoundValueList(printer, inputArgs, getInputs(), ListDelimiter::Paren);

  if (auto coreIdsAttr = (*this)->getAttrOfType<DenseI32ArrayAttr>(onnx_mlir::kCoreIdsAttrName))
    printCoreIdList(printer, "coreIds", coreIdsAttr.asArrayRef());
@@ -49,51 +184,57 @@ void PimCoreBatchOp::print(OpAsmPrinter& printer) {
  printer.printOptionalAttrDict(
    (*this)->getAttrs(),
    {getLaneCountAttrName().getValue(), getOperandSegmentSizesAttrName().getValue(), onnx_mlir::kCoreIdsAttrName});
-  printer << " ";
-  printer.printRegion(getBody(), /*printEntryBlockArgs=*/false);
  printer << " : ";
-  if (getLaneCount() > 1 && hasRepeatedTuple(TypeRange(getWeights()), weightsPerLane))
-    printTypeTupleRun(printer, TypeRange(getWeights()), weightsPerLane, ListDelimiter::Paren);
-  else
-    printCompressedTypeList(printer, TypeRange(getWeights()), ListDelimiter::Paren);
+  printCompressedTypeList(printer, TypeRange(getWeights()), ListDelimiter::Square);
  printer << " ";
-  printCompressedTypeList(printer, TypeRange(getInputs()), ListDelimiter::Square);
-  printer << " -> ()";
+  printCompressedTypeList(printer, TypeRange(getInputs()), ListDelimiter::Paren);
+  printer << " -> () ";
+  printer.printRegion(getBody(), /*printEntryBlockArgs=*/false);
 }

 ParseResult PimCoreBatchOp::parse(OpAsmParser& parser, OperationState& result) {
+  int64_t lowerBound = 0;
  int32_t laneCount = 0;
+  OpAsmParser::Argument laneArg;
+  SmallVector<OpAsmParser::Argument> weightArgs;
+  SmallVector<OpAsmParser::Argument> inputArgs;
+  SmallVector<OpAsmParser::Argument> regionArgs;
  SmallVector<OpAsmParser::UnresolvedOperand> weights;
  SmallVector<OpAsmParser::UnresolvedOperand> inputs;
  SmallVector<Type> weightTypes;
  SmallVector<Type> inputTypes;
  SmallVector<int32_t> coreIds;

-  if (parser.parseKeyword("lanes") || parser.parseInteger(laneCount)
-      || parseCompressedOrTupleOperandList(parser, ListDelimiter::Paren, weights)
-      || parseCompressedOperandList(parser, ListDelimiter::Square, inputs))
+  if (parser.parseArgument(laneArg) || parser.parseEqual() || parser.parseInteger(lowerBound)
+      || parser.parseKeyword("to") || parser.parseInteger(laneCount))
+    return failure();
+  if (lowerBound != 0)
+    return parser.emitError(parser.getCurrentLocation(), "core_batch currently requires a zero lower bound");
+
+  if (parseBoundValueList(parser, ListDelimiter::Square, weightArgs, weights)
+      || parseBoundValueList(parser, ListDelimiter::Paren, inputArgs, inputs))
    return failure();

-  bool hasCoreIds = succeeded(parser.parseOptionalKeyword("coreIds"));
+  bool hasCoreIds = parseOptionalKeywordAlias(parser, "coreIds", "core_ids");
  if (hasCoreIds && parseCompressedIntegerList(parser, coreIds))
    return failure();

-  if (parser.parseOptionalAttrDict(result.attributes))
-    return failure();
-
-  Region* body = result.addRegion();
-  if (parser.parseRegion(*body))
-    return failure();
-
-  if (parser.parseColon() || parseCompressedOrTupleTypeList(parser, ListDelimiter::Paren, weightTypes)
-      || parseCompressedTypeList(parser, ListDelimiter::Square, inputTypes) || parser.parseArrow()
-      || parser.parseLParen() || parser.parseRParen())
+  if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()
+      || parseCompressedRepeatedList(
+        parser, ListDelimiter::Square, weightTypes, [&](Type& type) { return parser.parseType(type); })
+      || parseCompressedRepeatedList(
+        parser, ListDelimiter::Paren, inputTypes, [&](Type& type) { return parser.parseType(type); })
+      || parser.parseArrow() || parser.parseLParen() || parser.parseRParen())
    return failure();

  if (weights.size() != weightTypes.size())
    return parser.emitError(parser.getCurrentLocation(), "number of weights and weight types must match");
+  if (weightArgs.size() != weights.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of weight bindings and weight operands must match");
  if (inputs.size() != inputTypes.size())
    return parser.emitError(parser.getCurrentLocation(), "number of inputs and input types must match");
+  if (inputArgs.size() != inputs.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of input bindings and input operands must match");
  if (hasCoreIds && result.attributes.get(onnx_mlir::kCoreIdsAttrName))
    return parser.emitError(parser.getCurrentLocation(),
                            "coreIds cannot be specified both positionally and in attr-dict");
@@ -110,7 +251,15 @@ ParseResult PimCoreBatchOp::parse(OpAsmParser& parser, OperationState& result) {
      || parser.resolveOperands(inputs, inputTypes, parser.getCurrentLocation(), result.operands)) {
    return failure();
  }
-  return success();
+
+  Region* body = result.addRegion();
+  laneArg.type = builder.getIndexType();
+  regionArgs.push_back(laneArg);
+  applyArgumentTypes(weightTypes, weightArgs);
+  llvm::append_range(regionArgs, weightArgs);
+  applyArgumentTypes(inputTypes, inputArgs);
+  llvm::append_range(regionArgs, inputArgs);
+  return parser.parseRegion(*body, regionArgs);
 }

 void PimYieldOp::print(OpAsmPrinter& printer) {