[AsmParser/Printer] Rework sourceloc support for function arguments.

When Location tracking support for block arguments was added, we
discussed various approaches to threading support for this through
function-like argument parsing.  At the time, we added a parallel array
of locations that could hold this.  It turns out that that approach was
verbose and error prone, roughly no one adopted it.

This patch takes a different approach, adding an optional source
locator to the UnresolvedOperand class.  This fits much more naturally
into the standard structure we use for representing locators, and gives
all the function like dialects locator support for free (e.g. see the
test adding an example for the LLVM dialect).

Differential Revision: https://reviews.llvm.org/D124188

mlir/include/mlir/IR/FunctionImplementation.h

@@ -60,7 +60,7 @@ ParseResult parseFunctionArgumentList(
     OpAsmParser &parser, bool allowAttributes, bool allowVariadic,
     SmallVectorImpl<OpAsmParser::UnresolvedOperand> &argNames,
     SmallVectorImpl<Type> &argTypes, SmallVectorImpl<NamedAttrList> &argAttrs,
-    SmallVectorImpl<Location> &argLocations, bool &isVariadic);
+    bool &isVariadic);
 
 /// Parses a function signature using `parser`. The `allowVariadic` argument
 /// indicates whether functions with variadic arguments are supported. The
@@ -70,8 +70,7 @@ ParseResult parseFunctionSignature(
     OpAsmParser &parser, bool allowVariadic,
     SmallVectorImpl<OpAsmParser::UnresolvedOperand> &argNames,
     SmallVectorImpl<Type> &argTypes, SmallVectorImpl<NamedAttrList> &argAttrs,
-    SmallVectorImpl<Location> &argLocations, bool &isVariadic,
+    bool &isVariadic, SmallVectorImpl<Type> &resultTypes,
-    SmallVectorImpl<Type> &resultTypes,
     SmallVectorImpl<NamedAttrList> &resultAttrs);
 
 /// Parser implementation for function-like operations.  Uses

mlir/include/mlir/IR/OpImplementation.h

@@ -1087,6 +1087,7 @@ public:
     SMLoc location;  // Location of the token.
     StringRef name;  // Value name, e.g. %42 or %abc
     unsigned number; // Number, e.g. 12 for an operand like %xyz#12
+    Optional<Location> sourceLoc; // Source location specifier if present.
   };
 
   /// Parse different components, viz., use-info of operand(s), successor(s),
@@ -1212,23 +1213,20 @@ public:
 
   /// Parses a region. Any parsed blocks are appended to 'region' and must be
   /// moved to the op regions after the op is created. The first block of the
-  /// region takes 'arguments' of types 'argTypes'. If `argLocations` is
+  /// region takes 'arguments' of types 'argTypes'.  If 'enableNameShadowing' is
-  /// non-empty it contains a location to be attached to each argument. If
+  /// set to true, the argument names are allowed to shadow the names of other
-  /// 'enableNameShadowing' is set to true, the argument names are allowed to
+  /// existing SSA values defined above the region scope. 'enableNameShadowing'
-  /// shadow the names of other existing SSA values defined above the region
+  /// can only be set to true for regions attached to operations that are
-  /// scope. 'enableNameShadowing' can only be set to true for regions attached
+  /// 'IsolatedFromAbove'.
-  /// to operations that are 'IsolatedFromAbove'.
   virtual ParseResult parseRegion(Region &region,
                                   ArrayRef<UnresolvedOperand> arguments = {},
                                   ArrayRef<Type> argTypes = {},
-                                  ArrayRef<Location> argLocations = {},
                                   bool enableNameShadowing = false) = 0;
 
   /// Parses a region if present.
   virtual OptionalParseResult parseOptionalRegion(
       Region &region, ArrayRef<UnresolvedOperand> arguments = {},
-      ArrayRef<Type> argTypes = {}, ArrayRef<Location> argLocations = {},
+      ArrayRef<Type> argTypes = {}, bool enableNameShadowing = false) = 0;
-      bool enableNameShadowing = false) = 0;
 
   /// Parses a region if present. If the region is present, a new region is
   /// allocated and placed in `region`. If no region is present or on failure,

mlir/lib/Dialect/Async/IR/Async.cpp

@@ -229,7 +229,6 @@ ParseResult ExecuteOp::parse(OpAsmParser &parser, OperationState &result) {
   Region *body = result.addRegion();
   if (parser.parseRegion(*body, /*arguments=*/{unwrappedArgs},
                          /*argTypes=*/{unwrappedTypes},
-                         /*argLocations=*/{},
                          /*enableNameShadowing=*/false))
     return failure();
 

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

@@ -759,12 +759,10 @@ static ParseResult parseLaunchFuncOperands(
   if (parser.parseOptionalKeyword("args"))
     return success();
   SmallVector<NamedAttrList> argAttrs;
-  SmallVector<Location> argLocations;
   bool isVariadic = false;
   return function_interface_impl::parseFunctionArgumentList(
       parser, /*allowAttributes=*/false,
-      /*allowVariadic=*/false, argNames, argTypes, argAttrs, argLocations,
+      /*allowVariadic=*/false, argNames, argTypes, argAttrs, isVariadic);
-      isVariadic);
 }
 
 static void printLaunchFuncOperands(OpAsmPrinter &printer, Operation *,
@@ -892,7 +890,6 @@ ParseResult GPUFuncOp::parse(OpAsmParser &parser, OperationState &result) {
   SmallVector<NamedAttrList> resultAttrs;
   SmallVector<Type> argTypes;
   SmallVector<Type> resultTypes;
-  SmallVector<Location> argLocations;
   bool isVariadic;
 
   // Parse the function name.
@@ -904,7 +901,7 @@ ParseResult GPUFuncOp::parse(OpAsmParser &parser, OperationState &result) {
   auto signatureLocation = parser.getCurrentLocation();
   if (failed(function_interface_impl::parseFunctionSignature(
           parser, /*allowVariadic=*/false, entryArgs, argTypes, argAttrs,
-          argLocations, isVariadic, resultTypes, resultAttrs)))
+          isVariadic, resultTypes, resultAttrs)))
     return failure();
 
   if (entryArgs.empty() && !argTypes.empty())

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

@@ -2151,7 +2151,6 @@ ParseResult LLVMFuncOp::parse(OpAsmParser &parser, OperationState &result) {
   SmallVector<NamedAttrList> resultAttrs;
   SmallVector<Type> argTypes;
   SmallVector<Type> resultTypes;
-  SmallVector<Location> argLocations;
   bool isVariadic;
 
   auto signatureLocation = parser.getCurrentLocation();
@@ -2159,7 +2158,7 @@ ParseResult LLVMFuncOp::parse(OpAsmParser &parser, OperationState &result) {
                              result.attributes) ||
       function_interface_impl::parseFunctionSignature(
           parser, /*allowVariadic=*/true, entryArgs, argTypes, argAttrs,
-          argLocations, isVariadic, resultTypes, resultAttrs))
+          isVariadic, resultTypes, resultAttrs))
     return failure();
 
   auto type =

mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp

@@ -2198,7 +2198,6 @@ ParseResult spirv::FuncOp::parse(OpAsmParser &parser, OperationState &state) {
   SmallVector<NamedAttrList> resultAttrs;
   SmallVector<Type> argTypes;
   SmallVector<Type> resultTypes;
-  SmallVector<Location> argLocations;
   auto &builder = parser.getBuilder();
 
   // Parse the name as a symbol.
@@ -2211,7 +2210,7 @@ ParseResult spirv::FuncOp::parse(OpAsmParser &parser, OperationState &state) {
   bool isVariadic = false;
   if (function_interface_impl::parseFunctionSignature(
           parser, /*allowVariadic=*/false, entryArgs, argTypes, argAttrs,
-          argLocations, isVariadic, resultTypes, resultAttrs))
+          isVariadic, resultTypes, resultAttrs))
     return failure();
 
   auto fnType = builder.getFunctionType(argTypes, resultTypes);

mlir/lib/IR/FunctionImplementation.cpp

@@ -17,7 +17,7 @@ ParseResult mlir::function_interface_impl::parseFunctionArgumentList(
     OpAsmParser &parser, bool allowAttributes, bool allowVariadic,
     SmallVectorImpl<OpAsmParser::UnresolvedOperand> &argNames,
     SmallVectorImpl<Type> &argTypes, SmallVectorImpl<NamedAttrList> &argAttrs,
-    SmallVectorImpl<Location> &argLocations, bool &isVariadic) {
+    bool &isVariadic) {
   if (parser.parseLParen())
     return failure();
 
@@ -35,8 +35,8 @@ ParseResult mlir::function_interface_impl::parseFunctionArgumentList(
       // Reject this if the preceding argument was missing a name.
       if (argNames.empty() && !argTypes.empty())
         return parser.emitError(loc, "expected type instead of SSA identifier");
-      argNames.push_back(argument);
 
+      // Parse required type.
       if (parser.parseColonType(argumentType))
         return failure();
     } else if (allowVariadic && succeeded(parser.parseOptionalEllipsis())) {
@@ -52,23 +52,19 @@ ParseResult mlir::function_interface_impl::parseFunctionArgumentList(
     // Add the argument type.
     argTypes.push_back(argumentType);
 
-    // Parse any argument attributes.
+    // Parse any argument attributes and source location information.
     NamedAttrList attrs;
-    if (parser.parseOptionalAttrDict(attrs))
+    if (parser.parseOptionalAttrDict(attrs) ||
+        parser.parseOptionalLocationSpecifier(argument.sourceLoc))
       return failure();
+         
     if (!allowAttributes && !attrs.empty())
       return parser.emitError(loc, "expected arguments without attributes");
     argAttrs.push_back(attrs);
 
-    // Parse a location if specified.
+    // If we had an argument name, then remember the parsed argument.
-    Optional<Location> explicitLoc;
+    if (!argument.name.empty())
-    if (!argument.name.empty() &&
+      argNames.push_back(argument);
-        parser.parseOptionalLocationSpecifier(explicitLoc))
-      return failure();
-    if (!explicitLoc)
-      explicitLoc = parser.getEncodedSourceLoc(loc);
-    argLocations.push_back(*explicitLoc);
-
     return success();
   };
 
@@ -135,12 +131,11 @@ ParseResult mlir::function_interface_impl::parseFunctionSignature(
     OpAsmParser &parser, bool allowVariadic,
     SmallVectorImpl<OpAsmParser::UnresolvedOperand> &argNames,
     SmallVectorImpl<Type> &argTypes, SmallVectorImpl<NamedAttrList> &argAttrs,
-    SmallVectorImpl<Location> &argLocations, bool &isVariadic,
+    bool &isVariadic, SmallVectorImpl<Type> &resultTypes,
-    SmallVectorImpl<Type> &resultTypes,
     SmallVectorImpl<NamedAttrList> &resultAttrs) {
   bool allowArgAttrs = true;
   if (parseFunctionArgumentList(parser, allowArgAttrs, allowVariadic, argNames,
-                                argTypes, argAttrs, argLocations, isVariadic))
+                                argTypes, argAttrs, isVariadic))
     return failure();
   if (succeeded(parser.parseOptionalArrow()))
     return parseFunctionResultList(parser, resultTypes, resultAttrs);
@@ -199,7 +194,6 @@ ParseResult mlir::function_interface_impl::parseFunctionOp(
   SmallVector<NamedAttrList> resultAttrs;
   SmallVector<Type> argTypes;
   SmallVector<Type> resultTypes;
-  SmallVector<Location> argLocations;
   auto &builder = parser.getBuilder();
 
   // Parse visibility.
@@ -215,8 +209,7 @@ ParseResult mlir::function_interface_impl::parseFunctionOp(
   SMLoc signatureLocation = parser.getCurrentLocation();
   bool isVariadic = false;
   if (parseFunctionSignature(parser, allowVariadic, entryArgs, argTypes,
-                             argAttrs, argLocations, isVariadic, resultTypes,
+                             argAttrs, isVariadic, resultTypes, resultAttrs))
-                             resultAttrs))
     return failure();
 
   std::string errorMessage;
@@ -259,7 +252,6 @@ ParseResult mlir::function_interface_impl::parseFunctionOp(
   SMLoc loc = parser.getCurrentLocation();
   OptionalParseResult parseResult = parser.parseOptionalRegion(
       *body, entryArgs, entryArgs.empty() ? ArrayRef<Type>() : argTypes,
-      entryArgs.empty() ? ArrayRef<Location>() : argLocations,
       /*enableNameShadowing=*/false);
   if (parseResult.hasValue()) {
     if (failed(*parseResult))

mlir/lib/Parser/Parser.cpp

@@ -360,20 +360,18 @@ public:
   //===--------------------------------------------------------------------===//
 
   /// Parse a region into 'region' with the provided entry block arguments.
-  /// If non-empty, 'argLocations' contains an optional locations for each
+  /// 'isIsolatedNameScope' indicates if the naming scope of this region is
-  /// argument. 'isIsolatedNameScope' indicates if the naming scope of this
+  /// isolated from those above.
-  /// region is isolated from those above.
   ParseResult
   parseRegion(Region &region,
               ArrayRef<std::pair<UnresolvedOperand, Type>> entryArguments,
-              ArrayRef<Location> argLocations,
               bool isIsolatedNameScope = false);
 
   /// Parse a region body into 'region'.
   ParseResult
   parseRegionBody(Region &region, SMLoc startLoc,
                   ArrayRef<std::pair<UnresolvedOperand, Type>> entryArguments,
-                  ArrayRef<Location> argLocations, bool isIsolatedNameScope);
+                  bool isIsolatedNameScope);
 
   //===--------------------------------------------------------------------===//
   // Block Parsing
@@ -943,7 +941,7 @@ ParseResult OperationParser::parseOperation() {
     unsigned opResI = 0;
     for (ResultRecord &resIt : resultIDs) {
       for (unsigned subRes : llvm::seq<unsigned>(0, std::get<1>(resIt))) {
-        if (addDefinition({std::get<2>(resIt), std::get<0>(resIt), subRes},
+        if (addDefinition({std::get<2>(resIt), std::get<0>(resIt), subRes, {}},
                           op->getResult(opResI++)))
           return failure();
       }
@@ -1049,8 +1047,7 @@ ParseResult OperationParser::parseGenericOperationAfterOpName(
       do {
         // Create temporary regions with the top level region as parent.
         result.regions.emplace_back(new Region(topLevelOp));
-        if (parseRegion(*result.regions.back(), /*entryArguments=*/{},
+        if (parseRegion(*result.regions.back(), /*entryArguments=*/{}))
-                        /*argLocations=*/{}))
           return failure();
       } while (consumeIf(Token::comma));
       if (parseToken(Token::r_paren, "expected ')' to end region list"))
@@ -1275,8 +1272,10 @@ public:
     if (parser.parseSSAUse(useInfo))
       return failure();
 
-    result = {useInfo.location, useInfo.name, useInfo.number};
+    result = {useInfo.location, useInfo.name, useInfo.number, {}};
-    return success();
+
+    // Parse a source locator on the operand if present.
+    return parseOptionalLocationSpecifier(result.sourceLoc);
   }
 
   /// Parse a single operand if present.
@@ -1429,7 +1428,6 @@ public:
   /// effectively defines the SSA values of `arguments` and assigns their type.
   ParseResult parseRegion(Region &region, ArrayRef<UnresolvedOperand> arguments,
                           ArrayRef<Type> argTypes,
-                          ArrayRef<Location> argLocations,
                           bool enableNameShadowing) override {
     assert(arguments.size() == argTypes.size() &&
            "mismatching number of arguments and types");
@@ -1443,8 +1441,7 @@ public:
     (void)isIsolatedFromAbove;
     assert((!enableNameShadowing || isIsolatedFromAbove) &&
            "name shadowing is only allowed on isolated regions");
-    if (parser.parseRegion(region, regionArguments, argLocations,
+    if (parser.parseRegion(region, regionArguments, enableNameShadowing))
-                           enableNameShadowing))
       return failure();
     return success();
   }
@@ -1453,12 +1450,10 @@ public:
   OptionalParseResult parseOptionalRegion(Region &region,
                                           ArrayRef<UnresolvedOperand> arguments,
                                           ArrayRef<Type> argTypes,
-                                          ArrayRef<Location> argLocations,
                                           bool enableNameShadowing) override {
     if (parser.getToken().isNot(Token::l_brace))
       return llvm::None;
-    return parseRegion(region, arguments, argTypes, argLocations,
+    return parseRegion(region, arguments, argTypes, enableNameShadowing);
-                       enableNameShadowing);
   }
 
   /// Parses a region if present. If the region is present, a new region is
@@ -1470,8 +1465,7 @@ public:
     if (parser.getToken().isNot(Token::l_brace))
       return llvm::None;
     std::unique_ptr<Region> newRegion = std::make_unique<Region>();
-    if (parseRegion(*newRegion, arguments, argTypes, /*argLocations=*/{},
+    if (parseRegion(*newRegion, arguments, argTypes, enableNameShadowing))
-                    enableNameShadowing))
       return failure();
 
     region = std::move(newRegion);
@@ -1798,7 +1792,7 @@ ParseResult OperationParser::parseRegion(
     Region &region,
     ArrayRef<std::pair<OperationParser::UnresolvedOperand, Type>>
         entryArguments,
-    ArrayRef<Location> argLocations, bool isIsolatedNameScope) {
+    bool isIsolatedNameScope) {
   // Parse the '{'.
   Token lBraceTok = getToken();
   if (parseToken(Token::l_brace, "expected '{' to begin a region"))
@@ -1810,7 +1804,7 @@ ParseResult OperationParser::parseRegion(
 
   // Parse the region body.
   if ((!entryArguments.empty() || getToken().isNot(Token::r_brace)) &&
-      parseRegionBody(region, lBraceTok.getLoc(), entryArguments, argLocations,
+      parseRegionBody(region, lBraceTok.getLoc(), entryArguments,
                       isIsolatedNameScope)) {
     return failure();
   }
@@ -1827,8 +1821,7 @@ ParseResult OperationParser::parseRegionBody(
     Region &region, SMLoc startLoc,
     ArrayRef<std::pair<OperationParser::UnresolvedOperand, Type>>
         entryArguments,
-    ArrayRef<Location> argLocations, bool isIsolatedNameScope) {
+    bool isIsolatedNameScope) {
-  assert(argLocations.empty() || argLocations.size() == entryArguments.size());
   auto currentPt = opBuilder.saveInsertionPoint();
 
   // Push a new named value scope.
@@ -1850,9 +1843,7 @@ ParseResult OperationParser::parseRegionBody(
     if (getToken().is(Token::caret_identifier))
       return emitError("invalid block name in region with named arguments");
 
-    for (const auto &it : llvm::enumerate(entryArguments)) {
+    for (auto &placeholderArgPair : entryArguments) {
-      size_t argIndex = it.index();
-      auto &placeholderArgPair = it.value();
       auto &argInfo = placeholderArgPair.first;
 
       // Ensure that the argument was not already defined.
@@ -1863,11 +1854,10 @@ ParseResult OperationParser::parseRegionBody(
                    .attachNote(getEncodedSourceLocation(*defLoc))
                << "previously referenced here";
       }
-      BlockArgument arg = block->addArgument(
+      Location loc = argInfo.sourceLoc.hasValue()
-          placeholderArgPair.second,
+                         ? argInfo.sourceLoc.getValue()
-          argLocations.empty()
+                         : getEncodedSourceLocation(argInfo.location);
-              ? getEncodedSourceLocation(placeholderArgPair.first.location)
+      BlockArgument arg = block->addArgument(placeholderArgPair.second, loc);
-              : argLocations[argIndex]);
 
       // Add a definition of this arg to the assembly state if provided.
       if (state.asmState)

mlir/test/Dialect/LLVMIR/func.mlir

@@ -1,5 +1,6 @@
 // RUN: mlir-opt -split-input-file -verify-diagnostics %s | mlir-opt | FileCheck %s
 // RUN: mlir-opt -split-input-file -verify-diagnostics -mlir-print-op-generic %s | FileCheck %s --check-prefix=GENERIC
+// RUN: mlir-opt -split-input-file -verify-diagnostics %s -mlir-print-debuginfo | mlir-opt -mlir-print-debuginfo | FileCheck %s --check-prefix=LOCINFO
 
 module {
   // GENERIC: "llvm.func"
@@ -93,7 +94,8 @@ module {
   }
 
   // CHECK: llvm.func @sretattr(%{{.*}}: !llvm.ptr<i32> {llvm.sret})
-  llvm.func @sretattr(%arg0: !llvm.ptr<i32> {llvm.sret}) {
+  // LOCINFO: llvm.func @sretattr(%{{.*}}: !llvm.ptr<i32> {llvm.sret} loc("some_source_loc"))
+  llvm.func @sretattr(%arg0: !llvm.ptr<i32> {llvm.sret} loc("some_source_loc")) {
     llvm.return
   }
 

mlir/test/lib/Dialect/Test/TestDialect.cpp

@@ -643,7 +643,7 @@ ParseResult IsolatedRegionOp::parse(OpAsmParser &parser,
 
   // Parse the body region, and reuse the operand info as the argument info.
   Region *body = result.addRegion();
-  return parser.parseRegion(*body, argInfo, argType, /*argLocations=*/{},
+  return parser.parseRegion(*body, argInfo, argType,
                             /*enableNameShadowing=*/true);
 }