407 lines · cpp
1//===- TestFormatUtils.cpp - MLIR Test Dialect Assembly Format Utilities --===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "TestFormatUtils.h"10#include "mlir/IR/Builders.h"11 12using namespace mlir;13using namespace test;14 15//===----------------------------------------------------------------------===//16// CustomDirectiveOperands17//===----------------------------------------------------------------------===//18 19ParseResult test::parseCustomDirectiveOperands(20 OpAsmParser &parser, OpAsmParser::UnresolvedOperand &operand,21 std::optional<OpAsmParser::UnresolvedOperand> &optOperand,22 SmallVectorImpl<OpAsmParser::UnresolvedOperand> &varOperands) {23 if (parser.parseOperand(operand))24 return failure();25 if (succeeded(parser.parseOptionalComma())) {26 optOperand.emplace();27 if (parser.parseOperand(*optOperand))28 return failure();29 }30 if (parser.parseArrow() || parser.parseLParen() ||31 parser.parseOperandList(varOperands) || parser.parseRParen())32 return failure();33 return success();34}35 36void test::printCustomDirectiveOperands(OpAsmPrinter &printer, Operation *,37 Value operand, Value optOperand,38 OperandRange varOperands) {39 printer << operand;40 if (optOperand)41 printer << ", " << optOperand;42 printer << " -> (" << varOperands << ")";43}44 45//===----------------------------------------------------------------------===//46// CustomDirectiveResults47//===----------------------------------------------------------------------===//48 49ParseResult50test::parseCustomDirectiveResults(OpAsmParser &parser, Type &operandType,51 Type &optOperandType,52 SmallVectorImpl<Type> &varOperandTypes) {53 if (parser.parseColon())54 return failure();55 56 if (parser.parseType(operandType))57 return failure();58 if (succeeded(parser.parseOptionalComma()))59 if (parser.parseType(optOperandType))60 return failure();61 if (parser.parseArrow() || parser.parseLParen() ||62 parser.parseTypeList(varOperandTypes) || parser.parseRParen())63 return failure();64 return success();65}66 67void test::printCustomDirectiveResults(OpAsmPrinter &printer, Operation *,68 Type operandType, Type optOperandType,69 TypeRange varOperandTypes) {70 printer << " : " << operandType;71 if (optOperandType)72 printer << ", " << optOperandType;73 printer << " -> (" << varOperandTypes << ")";74}75 76//===----------------------------------------------------------------------===//77// CustomDirectiveWithTypeRefs78//===----------------------------------------------------------------------===//79 80ParseResult test::parseCustomDirectiveWithTypeRefs(81 OpAsmParser &parser, Type operandType, Type optOperandType,82 const SmallVectorImpl<Type> &varOperandTypes) {83 if (parser.parseKeyword("type_refs_capture"))84 return failure();85 86 Type operandType2, optOperandType2;87 SmallVector<Type, 1> varOperandTypes2;88 if (parseCustomDirectiveResults(parser, operandType2, optOperandType2,89 varOperandTypes2))90 return failure();91 92 if (operandType != operandType2 || optOperandType != optOperandType2 ||93 varOperandTypes != varOperandTypes2)94 return failure();95 96 return success();97}98 99void test::printCustomDirectiveWithTypeRefs(OpAsmPrinter &printer,100 Operation *op, Type operandType,101 Type optOperandType,102 TypeRange varOperandTypes) {103 printer << " type_refs_capture ";104 printCustomDirectiveResults(printer, op, operandType, optOperandType,105 varOperandTypes);106}107 108//===----------------------------------------------------------------------===//109// CustomDirectiveOperandsAndTypes110//===----------------------------------------------------------------------===//111 112ParseResult test::parseCustomDirectiveOperandsAndTypes(113 OpAsmParser &parser, OpAsmParser::UnresolvedOperand &operand,114 std::optional<OpAsmParser::UnresolvedOperand> &optOperand,115 SmallVectorImpl<OpAsmParser::UnresolvedOperand> &varOperands,116 Type &operandType, Type &optOperandType,117 SmallVectorImpl<Type> &varOperandTypes) {118 if (parseCustomDirectiveOperands(parser, operand, optOperand, varOperands) ||119 parseCustomDirectiveResults(parser, operandType, optOperandType,120 varOperandTypes))121 return failure();122 return success();123}124 125void test::printCustomDirectiveOperandsAndTypes(126 OpAsmPrinter &printer, Operation *op, Value operand, Value optOperand,127 OperandRange varOperands, Type operandType, Type optOperandType,128 TypeRange varOperandTypes) {129 printCustomDirectiveOperands(printer, op, operand, optOperand, varOperands);130 printCustomDirectiveResults(printer, op, operandType, optOperandType,131 varOperandTypes);132}133 134//===----------------------------------------------------------------------===//135// CustomDirectiveRegions136//===----------------------------------------------------------------------===//137 138ParseResult test::parseCustomDirectiveRegions(139 OpAsmParser &parser, Region ®ion,140 SmallVectorImpl<std::unique_ptr<Region>> &varRegions) {141 if (parser.parseRegion(region))142 return failure();143 if (failed(parser.parseOptionalComma()))144 return success();145 std::unique_ptr<Region> varRegion = std::make_unique<Region>();146 if (parser.parseRegion(*varRegion))147 return failure();148 varRegions.emplace_back(std::move(varRegion));149 return success();150}151 152void test::printCustomDirectiveRegions(OpAsmPrinter &printer, Operation *,153 Region ®ion,154 MutableArrayRef<Region> varRegions) {155 printer.printRegion(region);156 if (!varRegions.empty()) {157 printer << ", ";158 for (Region ®ion : varRegions)159 printer.printRegion(region);160 }161}162 163//===----------------------------------------------------------------------===//164// CustomDirectiveSuccessors165//===----------------------------------------------------------------------===//166 167ParseResult168test::parseCustomDirectiveSuccessors(OpAsmParser &parser, Block *&successor,169 SmallVectorImpl<Block *> &varSuccessors) {170 if (parser.parseSuccessor(successor))171 return failure();172 if (failed(parser.parseOptionalComma()))173 return success();174 Block *varSuccessor;175 if (parser.parseSuccessor(varSuccessor))176 return failure();177 varSuccessors.append(2, varSuccessor);178 return success();179}180 181void test::printCustomDirectiveSuccessors(OpAsmPrinter &printer, Operation *,182 Block *successor,183 SuccessorRange varSuccessors) {184 printer << successor;185 if (!varSuccessors.empty())186 printer << ", " << varSuccessors.front();187}188 189//===----------------------------------------------------------------------===//190// CustomDirectiveAttributes191//===----------------------------------------------------------------------===//192 193ParseResult test::parseCustomDirectiveAttributes(OpAsmParser &parser,194 IntegerAttr &attr,195 IntegerAttr &optAttr) {196 if (parser.parseAttribute(attr))197 return failure();198 if (succeeded(parser.parseOptionalComma())) {199 if (parser.parseAttribute(optAttr))200 return failure();201 }202 return success();203}204 205void test::printCustomDirectiveAttributes(OpAsmPrinter &printer, Operation *,206 Attribute attribute,207 Attribute optAttribute) {208 printer << attribute;209 if (optAttribute)210 printer << ", " << optAttribute;211}212 213//===----------------------------------------------------------------------===//214// CustomDirectiveAttrDict215//===----------------------------------------------------------------------===//216 217ParseResult test::parseCustomDirectiveAttrDict(OpAsmParser &parser,218 NamedAttrList &attrs) {219 return parser.parseOptionalAttrDict(attrs);220}221 222void test::printCustomDirectiveAttrDict(OpAsmPrinter &printer, Operation *op,223 DictionaryAttr attrs) {224 printer.printOptionalAttrDict(attrs.getValue());225}226 227//===----------------------------------------------------------------------===//228// CustomDirectiveOptionalOperandRef229//===----------------------------------------------------------------------===//230 231ParseResult test::parseCustomDirectiveOptionalOperandRef(232 OpAsmParser &parser,233 std::optional<OpAsmParser::UnresolvedOperand> &optOperand) {234 int64_t operandCount = 0;235 if (parser.parseInteger(operandCount))236 return failure();237 bool expectedOptionalOperand = operandCount == 0;238 return success(expectedOptionalOperand != !!optOperand);239}240 241void test::printCustomDirectiveOptionalOperandRef(OpAsmPrinter &printer,242 Operation *op,243 Value optOperand) {244 printer << (optOperand ? "1" : "0");245}246 247//===----------------------------------------------------------------------===//248// CustomDirectiveOptionalOperand249//===----------------------------------------------------------------------===//250 251ParseResult test::parseCustomOptionalOperand(252 OpAsmParser &parser,253 std::optional<OpAsmParser::UnresolvedOperand> &optOperand) {254 if (succeeded(parser.parseOptionalLParen())) {255 optOperand.emplace();256 if (parser.parseOperand(*optOperand) || parser.parseRParen())257 return failure();258 }259 return success();260}261 262void test::printCustomOptionalOperand(OpAsmPrinter &printer, Operation *,263 Value optOperand) {264 if (optOperand)265 printer << "(" << optOperand << ") ";266}267 268//===----------------------------------------------------------------------===//269// CustomDirectiveSwitchCases270//===----------------------------------------------------------------------===//271 272ParseResult273test::parseSwitchCases(OpAsmParser &p, DenseI64ArrayAttr &cases,274 SmallVectorImpl<std::unique_ptr<Region>> &caseRegions) {275 SmallVector<int64_t> caseValues;276 while (succeeded(p.parseOptionalKeyword("case"))) {277 int64_t value;278 Region ®ion = *caseRegions.emplace_back(std::make_unique<Region>());279 if (p.parseInteger(value) || p.parseRegion(region, /*arguments=*/{}))280 return failure();281 caseValues.push_back(value);282 }283 cases = p.getBuilder().getDenseI64ArrayAttr(caseValues);284 return success();285}286 287void test::printSwitchCases(OpAsmPrinter &p, Operation *op,288 DenseI64ArrayAttr cases, RegionRange caseRegions) {289 for (auto [value, region] : llvm::zip(cases.asArrayRef(), caseRegions)) {290 p.printNewline();291 p << "case " << value << ' ';292 p.printRegion(*region, /*printEntryBlockArgs=*/false);293 }294}295 296//===----------------------------------------------------------------------===//297// CustomUsingPropertyInCustom298//===----------------------------------------------------------------------===//299 300bool test::parseUsingPropertyInCustom(OpAsmParser &parser,301 SmallVector<int64_t> &value) {302 auto elemParser = [&]() {303 int64_t v = 0;304 if (failed(parser.parseInteger(v)))305 return failure();306 value.push_back(v);307 return success();308 };309 return failed(parser.parseCommaSeparatedList(OpAsmParser::Delimiter::Square,310 elemParser));311}312 313void test::printUsingPropertyInCustom(OpAsmPrinter &printer, Operation *op,314 ArrayRef<int64_t> value) {315 printer << '[' << value << ']';316}317 318//===----------------------------------------------------------------------===//319// CustomDirectiveIntProperty320//===----------------------------------------------------------------------===//321 322bool test::parseIntProperty(OpAsmParser &parser, int64_t &value) {323 return failed(parser.parseInteger(value));324}325 326void test::printIntProperty(OpAsmPrinter &printer, Operation *op,327 int64_t value) {328 printer << value;329}330 331//===----------------------------------------------------------------------===//332// CustomDirectiveSumProperty333//===----------------------------------------------------------------------===//334 335bool test::parseSumProperty(OpAsmParser &parser, int64_t &second,336 int64_t first) {337 int64_t sum;338 auto loc = parser.getCurrentLocation();339 if (parser.parseInteger(second) || parser.parseEqual() ||340 parser.parseInteger(sum))341 return true;342 if (sum != second + first) {343 parser.emitError(loc, "Expected sum to equal first + second");344 return true;345 }346 return false;347}348 349void test::printSumProperty(OpAsmPrinter &printer, Operation *op,350 int64_t second, int64_t first) {351 printer << second << " = " << (second + first);352}353 354//===----------------------------------------------------------------------===//355// CustomDirectiveOptionalCustomParser356//===----------------------------------------------------------------------===//357 358OptionalParseResult test::parseOptionalCustomParser(AsmParser &p,359 IntegerAttr &result) {360 if (succeeded(p.parseOptionalKeyword("foo")))361 return p.parseAttribute(result);362 return {};363}364 365void test::printOptionalCustomParser(AsmPrinter &p, Operation *,366 IntegerAttr result) {367 p << "foo ";368 p.printAttribute(result);369}370 371//===----------------------------------------------------------------------===//372// CustomDirectiveAttrElideType373//===----------------------------------------------------------------------===//374 375ParseResult test::parseAttrElideType(AsmParser &parser, TypeAttr type,376 Attribute &attr) {377 return parser.parseAttribute(attr, type.getValue());378}379 380void test::printAttrElideType(AsmPrinter &printer, Operation *op, TypeAttr type,381 Attribute attr) {382 printer.printAttributeWithoutType(attr);383}384 385//===----------------------------------------------------------------------===//386// CustomDirectiveDummyRegionRef387//===----------------------------------------------------------------------===//388 389ParseResult test::parseDummyRegionRef(OpAsmParser &parser, Region ®ion) {390 return success();391}392 393void test::printDummyRegionRef(OpAsmPrinter &printer, Operation *op,394 Region ®ion) { /* do nothing */ }395 396//===----------------------------------------------------------------------===//397// CustomDirectiveDummySuccessorRef398//===----------------------------------------------------------------------===//399 400ParseResult test::parseDummySuccessorRef(OpAsmParser &parser,401 Block *successor) {402 return success();403}404 405void test::printDummySuccessorRef(OpAsmPrinter &printer, Operation *op,406 Block *successor) { /* do nothing */ }407