diff --git a/mlir/include/mlir/IR/OpImplementation.h b/mlir/include/mlir/IR/OpImplementation.h index accaa5fb46d339..954a9286266011 100644 --- a/mlir/include/mlir/IR/OpImplementation.h +++ b/mlir/include/mlir/IR/OpImplementation.h @@ -850,10 +850,6 @@ class AsmParser { StringRef attrName, NamedAttrList &attrs) = 0; - /// Parse a loc(...) specifier if present, filling in result if so. - virtual ParseResult - parseOptionalLocationSpecifier(Optional &result) = 0; - //===--------------------------------------------------------------------===// // Type Parsing //===--------------------------------------------------------------------===// @@ -1041,6 +1037,13 @@ class OpAsmParser : public AsmParser { using AsmParser::AsmParser; ~OpAsmParser() override; + /// Parse a loc(...) specifier if present, filling in result if so. + /// Location for BlockArgument and Operation may be deferred with an alias, in + /// which case an OpaqueLoc is set and will be resolved when parsing + /// completes. + virtual ParseResult + parseOptionalLocationSpecifier(Optional &result) = 0; + /// Return the name of the specified result in the specified syntax, as well /// as the sub-element in the name. It returns an empty string and ~0U for /// invalid result numbers. For example, in this operation: diff --git a/mlir/lib/Parser/AsmParserImpl.h b/mlir/lib/Parser/AsmParserImpl.h index 7e1095aae4b15d..2292060f35d447 100644 --- a/mlir/lib/Parser/AsmParserImpl.h +++ b/mlir/lib/Parser/AsmParserImpl.h @@ -429,22 +429,6 @@ class AsmParserImpl : public BaseT { return success(); } - /// Parse a loc(...) specifier if present, filling in result if so. - ParseResult - parseOptionalLocationSpecifier(Optional &result) override { - // If there is a 'loc' we parse a trailing location. - if (!parser.consumeIf(Token::kw_loc)) - return success(); - LocationAttr directLoc; - if (parser.parseToken(Token::l_paren, "expected '(' in location") || - parser.parseLocationInstance(directLoc) || - parser.parseToken(Token::r_paren, "expected ')' in location")) - return failure(); - - result = directLoc; - return success(); - } - //===--------------------------------------------------------------------===// // Type Parsing //===--------------------------------------------------------------------===// diff --git a/mlir/lib/Parser/Parser.cpp b/mlir/lib/Parser/Parser.cpp index e8a99f3b2bc78d..c97aad45a0744e 100644 --- a/mlir/lib/Parser/Parser.cpp +++ b/mlir/lib/Parser/Parser.cpp @@ -341,6 +341,11 @@ class OperationParser : public Parser { /// ParseResult parseTrailingLocationSpecifier(OpOrArgument opOrArgument); + /// Parse a location alias, that is a sequence looking like: #loc42 + /// The alias may have already be defined or may be defined later, in which + /// case an OpaqueLoc is used a placeholder. + ParseResult parseLocationAlias(LocationAttr &loc); + /// This is the structure of a result specifier in the assembly syntax, /// including the name, number of results, and location. using ResultRecord = std::tuple; @@ -1592,6 +1597,34 @@ class CustomOpAsmParser : public AsmParserImpl { return parser.parseCommaSeparatedListUntil(Token::r_paren, parseElt); } + /// Parse a loc(...) specifier if present, filling in result if so. + ParseResult + parseOptionalLocationSpecifier(Optional &result) override { + // If there is a 'loc' we parse a trailing location. + if (!parser.consumeIf(Token::kw_loc)) + return success(); + LocationAttr directLoc; + if (parser.parseToken(Token::l_paren, "expected '(' in location")) + return failure(); + + Token tok = parser.getToken(); + + // Check to see if we are parsing a location alias. + // Otherwise, we parse the location directly. + if (tok.is(Token::hash_identifier)) { + if (parser.parseLocationAlias(directLoc)) + return failure(); + } else if (parser.parseLocationInstance(directLoc)) { + return failure(); + } + + if (parser.parseToken(Token::r_paren, "expected ')' in location")) + return failure(); + + result = directLoc; + return success(); + } + private: /// Information about the result name specifiers. ArrayRef resultIDs; @@ -1719,6 +1752,33 @@ OperationParser::parseCustomOperation(ArrayRef resultIDs) { return op; } +ParseResult OperationParser::parseLocationAlias(LocationAttr &loc) { + Token tok = getToken(); + consumeToken(Token::hash_identifier); + StringRef identifier = tok.getSpelling().drop_front(); + if (identifier.contains('.')) { + return emitError(tok.getLoc()) + << "expected location, but found dialect attribute: '#" << identifier + << "'"; + } + + // If this alias can be resolved, do it now. + Attribute attr = state.symbols.attributeAliasDefinitions.lookup(identifier); + if (attr) { + if (!(loc = attr.dyn_cast())) + return emitError(tok.getLoc()) + << "expected location, but found '" << attr << "'"; + } else { + // Otherwise, remember this operation and resolve its location later. + // In the meantime, use a special OpaqueLoc as a marker. + loc = OpaqueLoc::get(deferredLocsReferences.size(), + TypeID::get(), + UnknownLoc::get(getContext())); + deferredLocsReferences.push_back(DeferredLocInfo{tok.getLoc(), identifier}); + } + return success(); +} + ParseResult OperationParser::parseTrailingLocationSpecifier(OpOrArgument opOrArgument) { // If there is a 'loc' we parse a trailing location. @@ -1729,34 +1789,11 @@ OperationParser::parseTrailingLocationSpecifier(OpOrArgument opOrArgument) { Token tok = getToken(); // Check to see if we are parsing a location alias. + // Otherwise, we parse the location directly. LocationAttr directLoc; if (tok.is(Token::hash_identifier)) { - consumeToken(); - - StringRef identifier = tok.getSpelling().drop_front(); - if (identifier.contains('.')) { - return emitError(tok.getLoc()) - << "expected location, but found dialect attribute: '#" - << identifier << "'"; - } - - // If this alias can be resolved, do it now. - Attribute attr = state.symbols.attributeAliasDefinitions.lookup(identifier); - if (attr) { - if (!(directLoc = attr.dyn_cast())) - return emitError(tok.getLoc()) - << "expected location, but found '" << attr << "'"; - } else { - // Otherwise, remember this operation and resolve its location later. - // In the meantime, use a special OpaqueLoc as a marker. - directLoc = OpaqueLoc::get(deferredLocsReferences.size(), - TypeID::get(), - UnknownLoc::get(getContext())); - deferredLocsReferences.push_back( - DeferredLocInfo{tok.getLoc(), identifier}); - } - - // Otherwise, we parse the location directly. + if (parseLocationAlias(directLoc)) + return failure(); } else if (parseLocationInstance(directLoc)) { return failure(); } diff --git a/mlir/test/IR/pretty_printed_region_op.mlir b/mlir/test/IR/pretty_printed_region_op.mlir index 3684acbe6eb487..7bdc06f581260c 100644 --- a/mlir/test/IR/pretty_printed_region_op.mlir +++ b/mlir/test/IR/pretty_printed_region_op.mlir @@ -1,5 +1,6 @@ -// RUN: mlir-opt -allow-unregistered-dialect -split-input-file %s | FileCheck %s --check-prefixes=CHECK-CUSTOM,CHECK +// RUN: mlir-opt -allow-unregistered-dialect -split-input-file %s | FileCheck %s --check-prefixes=CHECK-CUSTOM,CHECK // RUN: mlir-opt -allow-unregistered-dialect -mlir-print-op-generic -split-input-file %s | FileCheck %s --check-prefixes=CHECK,CHECK-GENERIC +// RUN: mlir-opt -allow-unregistered-dialect -split-input-file --mlir-print-op-generic --mlir-print-debuginfo -mlir-print-local-scope %s | FileCheck %s --check-prefixes=CHECK-LOCATION // ----- @@ -33,3 +34,16 @@ func @pretty_printed_region_op(%arg0 : f32, %arg1 : f32) -> (f32) { return %0 : f32 } +// ----- + + +func @pretty_printed_region_op_deferred_loc(%arg0 : f32, %arg1 : f32) -> (f32) { +// CHECK-LOCATION: "test.pretty_printed_region"(%arg1, %arg0) +// CHECK-LOCATION: ^bb0(%arg[[x:[0-9]+]]: f32 loc("foo"), %arg[[y:[0-9]+]]: f32 loc("foo") +// CHECK-LOCATION: %[[RES:.*]] = "special.op"(%arg[[x]], %arg[[y]]) : (f32, f32) -> f32 +// CHECK-LOCATION: "test.return"(%[[RES]]) : (f32) -> () +// CHECK-LOCATION: : (f32, f32) -> f32 + + %res = test.pretty_printed_region %arg1, %arg0 start special.op end : (f32, f32) -> (f32) loc("foo") + return %res : f32 +}