Preserve function argument locations.

[dominikgrewe]

Summary

Previously the optional locations of function arguments were dropped in parseFunctionArgumentList. This CL adds another output argument to the function through which they are now returned. The values are then plumbed through as an array of optional locations in the various places.

[jpienaar]

Nice!

mlir/lib/IR/FunctionImplementation.cpp

265

I'd expect one can just pass argLocations, as these shouldn't be populated if empty. What happens if you pass it?

mlir/lib/Parser/Parser.cpp

1820–1821

Could one use llvm::zip here to iterate over both together?

[dominikgrewe]

Updated function comments.

[dominikgrewe]

mlir/lib/IR/FunctionImplementation.cpp

265

I actually get quite a few test failures due to the assert I added to OperationParser::parseRegionBody.

mlir/lib/Parser/Parser.cpp

1820–1821

Not if argLocations is empty right?

[rriddle]

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

686–688

I'd just drop the explicit numbers now that SmallVector supports it.

mlir/lib/Parser/Parser.cpp

1820–1821

Nit: use llvm::enumerate() here.

1832–1834

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

598

Please add an argument comment here.

[dominikgrewe]

Addressed comments.

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

686–688

Changed it here and in other functions touched by this CL.

mlir/lib/Parser/Parser.cpp

1820–1821

Using llvm::enumerate now as per River's comment.

[dominikgrewe]

Addressed comments (with actual changes this time)

[rriddle]

LGTM, thanks for doing this! Really appreciated.

mlir/include/mlir/IR/FunctionImplementation.h

 /// function arguments and results, VariadicFlag indicates whether the function
 /// should have variadic arguments; in case of error, it may populate the last
 /// argument with a message.
 using FuncTypeBuilder = function_ref<Type(
     Builder &, ArrayRef<Type>, ArrayRef<Type>, VariadicFlag, std::string &)>;
 
 /// Parses function arguments using `parser`. The `allowVariadic` argument
 /// indicates whether functions with variadic arguments are supported. The
-/// trailing arguments are populated by this function with names, types and
-/// attributes of the arguments.
+/// trailing arguments are populated by this function with names, types,
+/// attributes and locations of the arguments.
 ParseResult parseFunctionArgumentList(
     OpAsmParser &parser, bool allowAttributes, bool allowVariadic,
     SmallVectorImpl<OpAsmParser::OperandType> &argNames,
     SmallVectorImpl<Type> &argTypes, SmallVectorImpl<NamedAttrList> &argAttrs,
-    bool &isVariadic);
+    SmallVectorImpl<Optional<Location>> &argLocations, bool &isVariadic);
 
 /// Parses a function signature using `parser`. The `allowVariadic` argument
 /// indicates whether functions with variadic arguments are supported. The
-/// trailing arguments are populated by this function with names, types and
-/// attributes of the arguments and those of the results.
+/// trailing arguments are populated by this function with names, types,
+/// attributes and locations of the arguments and those of the results.
 ParseResult
 parseFunctionSignature(OpAsmParser &parser, bool allowVariadic,
                        SmallVectorImpl<OpAsmParser::OperandType> &argNames,
                        SmallVectorImpl<Type> &argTypes,
                        SmallVectorImpl<NamedAttrList> &argAttrs,
+                       SmallVectorImpl<Optional<Location>> &argLocations,
                        bool &isVariadic, SmallVectorImpl<Type> &resultTypes,
                        SmallVectorImpl<NamedAttrList> &resultAttrs);
 
 /// Parser implementation for function-like operations.  Uses
 /// `funcTypeBuilder` to construct the custom function type given lists of
 /// input and output types.  If `allowVariadic` is set, the parser will accept
 /// trailing ellipsis in the function signature and indicate to the builder
 /// whether the function is variadic.  If the builder returns a null type,

mlir/include/mlir/IR/OpImplementation.h

parseAffineExprOfSSAIds(SmallVectorImpl<OperandType> &dimOperands,
                           AffineExpr &expr) = 0;
 
   //===--------------------------------------------------------------------===//
   // Region Parsing
   //===--------------------------------------------------------------------===//
 
   /// 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 'enableNameShadowing' is
-  /// set to true, the argument names are allowed to shadow the names of other
-  /// existing SSA values defined above the region scope. 'enableNameShadowing'
-  /// can only be set to true for regions attached to operations that are
-  /// 'IsolatedFromAbove.
-  virtual ParseResult parseRegion(Region &region,
-                                  ArrayRef<OperandType> arguments = {},
-                                  ArrayRef<Type> argTypes = {},
-                                  bool enableNameShadowing = false) = 0;
+  /// region takes 'arguments' of types 'argTypes'. If `argLocations` is
+  /// non-empty it contains an optional location to be attached to each
+  /// argument. If 'enableNameShadowing' is set to true, the argument names are
+  /// allowed to shadow the names of other existing SSA values defined above the
+  /// region scope. 'enableNameShadowing' can only be set to true for regions
+  /// attached to operations that are 'IsolatedFromAbove'.
+  virtual ParseResult
+  parseRegion(Region &region, ArrayRef<OperandType> arguments = {},
+              ArrayRef<Optional<Location>> argLocations = {},
 
   /// Parses a region if present.
   virtual OptionalParseResult
   parseOptionalRegion(Region &region, ArrayRef<OperandType> arguments = {},
                       ArrayRef<Type> argTypes = {},
+                      ArrayRef<Optional<Location>> argLocations = {},
                       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,
   /// `region` remains untouched.
   virtual OptionalParseResult parseOptionalRegion(
       std::unique_ptr<Region> &region, ArrayRef<OperandType> arguments = {},
       ArrayRef<Type> argTypes = {}, bool enableNameShadowing = false) = 0;

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

static ParseResult parseExecuteOp(OpAsmParser &parser, OperationState &result) {
   if (parser.parseOptionalAttrDictWithKeyword(attrs))
     return failure();
   result.addAttributes(attrs);
 
   // Parse asynchronous region.
   Region *body = result.addRegion();
   if (parser.parseRegion(*body, /*arguments=*/{unwrappedArgs},
                          /*argTypes=*/{unwrappedTypes},
+                         /*argLocations=*/{},
                          /*enableNameShadowing=*/false))
     return failure();
 
   return success();
 }
 
 static LogicalResult verify(ExecuteOp op) {
   // Unwrap async.execute value operands types.

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

 }
 
 static ParseResult
 parseLaunchFuncOperands(OpAsmParser &parser,
                         SmallVectorImpl<OpAsmParser::OperandType> &argNames,
                         SmallVectorImpl<Type> &argTypes) {
   if (parser.parseOptionalKeyword("args"))
     return success();
   SmallVector<NamedAttrList, 4> argAttrs;
+  SmallVector<Optional<Location>, 4> argLocations;
   bool isVariadic = false;
   return function_like_impl::parseFunctionArgumentList(
       parser, /*allowAttributes=*/false,
-      /*allowVariadic=*/false, argNames, argTypes, argAttrs, isVariadic);
+      /*allowVariadic=*/false, argNames, argTypes, argAttrs, argLocations,
+      isVariadic);
 }
 
 static void printLaunchFuncOperands(OpAsmPrinter &printer, Operation *,
                                     OperandRange operands, TypeRange types) {
   if (operands.empty())
     return;
   printer << "args(";
   llvm::interleaveComma(llvm::zip(operands, types), printer,

 ///                 (`->` function-result-list)? memory-attribution `kernel`?
 ///                 function-attributes? region
 static ParseResult parseGPUFuncOp(OpAsmParser &parser, OperationState &result) {
   SmallVector<OpAsmParser::OperandType, 8> entryArgs;
   SmallVector<NamedAttrList, 1> argAttrs;
   SmallVector<NamedAttrList, 1> resultAttrs;
   SmallVector<Type, 8> argTypes;
   SmallVector<Type, 4> resultTypes;
+  SmallVector<Optional<Location>, 8> argLocations;
   bool isVariadic;
 
   // Parse the function name.
   StringAttr nameAttr;
   if (parser.parseSymbolName(nameAttr, ::mlir::SymbolTable::getSymbolAttrName(),
                              result.attributes))
     return failure();
 
   auto signatureLocation = parser.getCurrentLocation();
   if (failed(function_like_impl::parseFunctionSignature(
           parser, /*allowVariadic=*/false, entryArgs, argTypes, argAttrs,
-          isVariadic, resultTypes, resultAttrs)))
+          argLocations, isVariadic, resultTypes, resultAttrs)))
     return failure();
 
   if (entryArgs.empty() && !argTypes.empty())
     return parser.emitError(signatureLocation)
            << "gpu.func requires named arguments";
 
   // Construct the function type. More types will be added to the region, but
   // not to the function type.

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

result.addAttribute(
                            parser, result, LLVM::Linkage::External)));
 
   StringAttr nameAttr;
   SmallVector<OpAsmParser::OperandType, 8> entryArgs;
   SmallVector<NamedAttrList, 1> argAttrs;
   SmallVector<NamedAttrList, 1> resultAttrs;
   SmallVector<Type, 8> argTypes;
   SmallVector<Type, 4> resultTypes;
+  SmallVector<Optional<Location>, 8> argLocations;
   bool isVariadic;
 
   auto signatureLocation = parser.getCurrentLocation();
   if (parser.parseSymbolName(nameAttr, SymbolTable::getSymbolAttrName(),
                              result.attributes) ||
       function_like_impl::parseFunctionSignature(
           parser, /*allowVariadic=*/true, entryArgs, argTypes, argAttrs,
-          isVariadic, resultTypes, resultAttrs))
+          argLocations, isVariadic, resultTypes, resultAttrs))
     return failure();
 
   auto type =
       buildLLVMFunctionType(parser, signatureLocation, argTypes, resultTypes,
                             function_like_impl::VariadicFlag(isVariadic));
   if (!type)
     return failure();
   result.addAttribute(function_like_impl::getTypeAttrName(),

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

 //===----------------------------------------------------------------------===//
 
 static ParseResult parseFuncOp(OpAsmParser &parser, OperationState &state) {
   SmallVector<OpAsmParser::OperandType, 4> entryArgs;
   SmallVector<NamedAttrList, 4> argAttrs;
   SmallVector<NamedAttrList, 4> resultAttrs;
   SmallVector<Type, 4> argTypes;
   SmallVector<Type, 4> resultTypes;
+  SmallVector<Optional<Location>, 4> argLocations;
   auto &builder = parser.getBuilder();
 
   // Parse the name as a symbol.
   StringAttr nameAttr;
   if (parser.parseSymbolName(nameAttr, SymbolTable::getSymbolAttrName(),
                              state.attributes))
     return failure();
 
   // Parse the function signature.
   bool isVariadic = false;
   if (function_like_impl::parseFunctionSignature(
           parser, /*allowVariadic=*/false, entryArgs, argTypes, argAttrs,
-          isVariadic, resultTypes, resultAttrs))
+          argLocations, isVariadic, resultTypes, resultAttrs))
     return failure();
 
   auto fnType = builder.getFunctionType(argTypes, resultTypes);
   state.addAttribute(function_like_impl::getTypeAttrName(),
                      TypeAttr::get(fnType));
 
   // Parse the optional function control keyword.
   spirv::FunctionControl fnControl;

mlir/lib/IR/FunctionImplementation.cpp

 #include "mlir/IR/SymbolTable.h"
 
 using namespace mlir;
 
 ParseResult mlir::function_like_impl::parseFunctionArgumentList(
     OpAsmParser &parser, bool allowAttributes, bool allowVariadic,
     SmallVectorImpl<OpAsmParser::OperandType> &argNames,
     SmallVectorImpl<Type> &argTypes, SmallVectorImpl<NamedAttrList> &argAttrs,
-    bool &isVariadic) {
+    SmallVectorImpl<Optional<Location>> &argLocations, bool &isVariadic) {
   if (parser.parseLParen())
     return failure();
 
   // The argument list either has to consistently have ssa-id's followed by
   // types, or just be a type list.  It isn't ok to sometimes have SSA ID's and
   // sometimes not.
   auto parseArgument = [&]() -> ParseResult {
     llvm::SMLoc loc = parser.getCurrentLocation();
auto parseArgument = [&]() -> ParseResult {
     // Parse any argument attributes.
     NamedAttrList attrs;
     if (parser.parseOptionalAttrDict(attrs))
       return failure();
     if (!allowAttributes && !attrs.empty())
       return parser.emitError(loc, "expected arguments without attributes");
     argAttrs.push_back(attrs);
 
-    // Parse a location if specified.  TODO: Don't drop it on the floor.
+    // Parse a location if specified.
     Optional<Location> explicitLoc;
     if (!argument.name.empty() &&
         parser.parseOptionalLocationSpecifier(explicitLoc))
       return failure();
+    argLocations.push_back(explicitLoc);
 
     return success();
   };
 
   // Parse the function arguments.
   isVariadic = false;
   if (failed(parser.parseOptionalRParen())) {
     do {
if (parser.parseType(resultTypes.back()) ||
       return failure();
     }
   } while (succeeded(parser.parseOptionalComma()));
   return parser.parseRParen();
 }
 
 /// Parses a function signature using `parser`. The `allowVariadic` argument
 /// indicates whether functions with variadic arguments are supported. The
-/// trailing arguments are populated by this function with names, types and
-/// attributes of the arguments and those of the results.
+/// trailing arguments are populated by this function with names, types,
+/// attributes and optional locations of the arguments and the types and
+/// attributes of the results.
 ParseResult mlir::function_like_impl::parseFunctionSignature(
     OpAsmParser &parser, bool allowVariadic,
     SmallVectorImpl<OpAsmParser::OperandType> &argNames,
     SmallVectorImpl<Type> &argTypes, SmallVectorImpl<NamedAttrList> &argAttrs,
-    bool &isVariadic, SmallVectorImpl<Type> &resultTypes,
+    SmallVectorImpl<Optional<Location>> &argLocations, bool &isVariadic,
+    SmallVectorImpl<Type> &resultTypes,
     SmallVectorImpl<NamedAttrList> &resultAttrs) {
   bool allowArgAttrs = true;
   if (parseFunctionArgumentList(parser, allowArgAttrs, allowVariadic, argNames,
-                                argTypes, argAttrs, isVariadic))
+                                argTypes, argAttrs, argLocations, isVariadic))
     return failure();
   if (succeeded(parser.parseOptionalArrow()))
     return parseFunctionResultList(parser, resultTypes, resultAttrs);
   return success();
 }
 
 /// Implementation of `addArgAndResultAttrs` that is attribute list type
 /// agnostic.

 ParseResult mlir::function_like_impl::parseFunctionLikeOp(
     OpAsmParser &parser, OperationState &result, bool allowVariadic,
     FuncTypeBuilder funcTypeBuilder) {
   SmallVector<OpAsmParser::OperandType, 4> entryArgs;
   SmallVector<NamedAttrList, 4> argAttrs;
   SmallVector<NamedAttrList, 4> resultAttrs;
   SmallVector<Type, 4> argTypes;
   SmallVector<Type, 4> resultTypes;
+  SmallVector<Optional<Location>, 4> argLocations;
   auto &builder = parser.getBuilder();
 
   // Parse visibility.
   impl::parseOptionalVisibilityKeyword(parser, result.attributes);
 
   // Parse the name as a symbol.
   StringAttr nameAttr;
   if (parser.parseSymbolName(nameAttr, SymbolTable::getSymbolAttrName(),
                              result.attributes))
     return failure();
 
   // Parse the function signature.
   llvm::SMLoc signatureLocation = parser.getCurrentLocation();
   bool isVariadic = false;
   if (parseFunctionSignature(parser, allowVariadic, entryArgs, argTypes,
-                             argAttrs, isVariadic, resultTypes, resultAttrs))
+                             argAttrs, argLocations, isVariadic, resultTypes,
+                             resultAttrs))
     return failure();
 
   std::string errorMessage;
   Type type = funcTypeBuilder(builder, argTypes, resultTypes,
                               VariadicFlag(isVariadic), errorMessage);
   if (!type) {
     return parser.emitError(signatureLocation)
            << "failed to construct function type"
ParseResult mlir::function_like_impl::parseFunctionLikeOp(
   addArgAndResultAttrs(builder, result, argAttrs, resultAttrs);
 
   // Parse the optional function body. The printer will not print the body if
   // its empty, so disallow parsing of empty body in the parser.
   auto *body = result.addRegion();
   llvm::SMLoc loc = parser.getCurrentLocation();
   OptionalParseResult parseResult = parser.parseOptionalRegion(
       *body, entryArgs, entryArgs.empty() ? ArrayRef<Type>() : argTypes,
+      entryArgs.empty() ? ArrayRef<Optional<Location>>() : argLocations,
       /*enableNameShadowing=*/false);
   if (parseResult.hasValue()) {
     if (failed(*parseResult))
       return failure();
     // Function body was parsed, make sure its not empty.
     if (body->empty())
       return parser.emitError(loc, "expected non-empty function body");
   }

mlir/lib/Parser/Parser.cpp

public:
   /// an object of type 'OperationName'. Otherwise, failure is returned.
   FailureOr<OperationName> parseCustomOperationName();
 
   //===--------------------------------------------------------------------===//
   // Region Parsing
   //===--------------------------------------------------------------------===//
 
   /// Parse a region into 'region' with the provided entry block arguments.
-  /// 'isIsolatedNameScope' indicates if the naming scope of this region is
-  /// isolated from those above.
+  /// If non-empty, 'argLocations' contains an optional locations for each
+  /// argument. 'isIsolatedNameScope' indicates if the naming scope of this
+  /// region is isolated from those above.
   ParseResult parseRegion(Region &region,
                           ArrayRef<std::pair<SSAUseInfo, Type>> entryArguments,
+                          ArrayRef<Optional<Location>> argLocations = {},
                           bool isIsolatedNameScope = false);
 
   /// Parse a region body into 'region'.
   ParseResult
   parseRegionBody(Region &region, llvm::SMLoc startLoc,
                   ArrayRef<std::pair<SSAUseInfo, Type>> entryArguments,
+                  ArrayRef<Optional<Location>> argLocations,
                   bool isIsolatedNameScope);
 
   //===--------------------------------------------------------------------===//
   // Block Parsing
   //===--------------------------------------------------------------------===//
 
   /// Parse a new block into 'block'.
   ParseResult parseBlock(Block *&block);
public:
   //===--------------------------------------------------------------------===//
   // Region Parsing
   //===--------------------------------------------------------------------===//
 
   /// Parse a region that takes `arguments` of `argTypes` types.  This
   /// effectively defines the SSA values of `arguments` and assigns their type.
   ParseResult parseRegion(Region &region, ArrayRef<OperandType> arguments,
                           ArrayRef<Type> argTypes,
+                          ArrayRef<Optional<Location>> argLocations,
                           bool enableNameShadowing) override {
     assert(arguments.size() == argTypes.size() &&
            "mismatching number of arguments and types");
 
     SmallVector<std::pair<OperationParser::SSAUseInfo, Type>, 2>
         regionArguments;
     for (auto pair : llvm::zip(arguments, argTypes)) {
       const OperandType &operand = std::get<0>(pair);
       Type type = std::get<1>(pair);
       OperationParser::SSAUseInfo operandInfo = {operand.name, operand.number,
                                                  operand.location};
       regionArguments.emplace_back(operandInfo, type);
     }
 
     // Try to parse the region.
     (void)isIsolatedFromAbove;
     assert((!enableNameShadowing || isIsolatedFromAbove) &&
            "name shadowing is only allowed on isolated regions");
-    if (parser.parseRegion(region, regionArguments, enableNameShadowing))
+    if (parser.parseRegion(region, regionArguments, argLocations,
+                           enableNameShadowing))
       return failure();
     return success();
   }
 
   /// Parses a region if present.
-  OptionalParseResult parseOptionalRegion(Region &region,
-                                          ArrayRef<OperandType> arguments,
-                                          ArrayRef<Type> argTypes,
-                                          bool enableNameShadowing) override {
+  OptionalParseResult
+  parseOptionalRegion(Region &region, ArrayRef<OperandType> arguments,
+                      ArrayRef<Optional<Location>> argLocations,
     if (parser.getToken().isNot(Token::l_brace))
       return llvm::None;
-    return parseRegion(region, arguments, argTypes, enableNameShadowing);
+    return parseRegion(region, arguments, argTypes, argLocations,
+                       enableNameShadowing);
   }
 
   /// Parses a region if present. If the region is present, a new region is
   /// allocated and placed in `region`. If no region is present, `region`
   /// remains untouched.
   OptionalParseResult
   parseOptionalRegion(std::unique_ptr<Region> &region,
                       ArrayRef<OperandType> arguments, ArrayRef<Type> argTypes,
                       bool enableNameShadowing = false) override {
     if (parser.getToken().isNot(Token::l_brace))
       return llvm::None;
     std::unique_ptr<Region> newRegion = std::make_unique<Region>();
-    if (parseRegion(*newRegion, arguments, argTypes, enableNameShadowing))
+    if (parseRegion(*newRegion, arguments, argTypes, /*argLocations=*/{},
+                    enableNameShadowing))
       return failure();
 
     region = std::move(newRegion);
     return success();
   }
 
   /// Parse a region argument. The type of the argument will be resolved later
   /// by a call to `parseRegion`.

 
 //===----------------------------------------------------------------------===//
 // Region Parsing
 //===----------------------------------------------------------------------===//
 
 ParseResult OperationParser::parseRegion(
     Region &region,
     ArrayRef<std::pair<OperationParser::SSAUseInfo, Type>> entryArguments,
-    bool isIsolatedNameScope) {
+    ArrayRef<Optional<Location>> argLocations, bool isIsolatedNameScope) {
   // Parse the '{'.
   Token lBraceTok = getToken();
   if (parseToken(Token::l_brace, "expected '{' to begin a region"))
     return failure();
 
   // If we are populating the parser state, start a new region definition.
   if (state.asmState)
     state.asmState->startRegionDefinition();
 
   // Parse the region body.
   if ((!entryArguments.empty() || getToken().isNot(Token::r_brace)) &&
-      parseRegionBody(region, lBraceTok.getLoc(), entryArguments,
+      parseRegionBody(region, lBraceTok.getLoc(), entryArguments, argLocations,
                       isIsolatedNameScope)) {
     return failure();
   }
   consumeToken(Token::r_brace);
 
   // If we are populating the parser state, finalize this region.
   if (state.asmState)
     state.asmState->finalizeRegionDefinition();
 
   return success();
 }
 
 ParseResult OperationParser::parseRegionBody(
     Region &region, llvm::SMLoc startLoc,
     ArrayRef<std::pair<OperationParser::SSAUseInfo, Type>> entryArguments,
-    bool isIsolatedNameScope) {
+    ArrayRef<Optional<Location>> argLocations, bool isIsolatedNameScope) {
+  assert(argLocations.empty() || argLocations.size() == entryArguments.size());
   auto currentPt = opBuilder.saveInsertionPoint();
 
   // Push a new named value scope.
   pushSSANameScope(isIsolatedNameScope);
 
   // Parse the first block directly to allow for it to be unnamed.
   auto owningBlock = std::make_unique<Block>();
   Block *block = owningBlock.get();
 
   // If this block is not defined in the source file, add a definition for it
   // now in the assembly state. Blocks with a name will be defined when the name
   // is parsed.
   if (state.asmState && getToken().isNot(Token::caret_identifier))
     state.asmState->addDefinition(block, startLoc);
 
   // Add arguments to the entry block.
   if (!entryArguments.empty()) {
     // If we had named arguments, then don't allow a block name.
     if (getToken().is(Token::caret_identifier))
       return emitError("invalid block name in region with named arguments");
 
-    for (auto &placeholderArgPair : entryArguments) {
+    for (int argIndex = 0; argIndex < entryArguments.size(); ++argIndex) {
+      auto &placeholderArgPair = entryArguments[argIndex];
       auto &argInfo = placeholderArgPair.first;
 
       // Ensure that the argument was not already defined.
       if (auto defLoc = getReferenceLoc(argInfo.name, argInfo.number)) {
         return emitError(argInfo.loc, "region entry argument '" + argInfo.name +
                                           "' is already in use")
                    .attachNote(getEncodedSourceLocation(*defLoc))
                << "previously referenced here";
       }
-      auto loc = getEncodedSourceLocation(placeholderArgPair.first.loc);
+      Location loc =
+          (!argLocations.empty() && argLocations[argIndex].hasValue())
+              ? argLocations[argIndex].getValue()
+              : getEncodedSourceLocation(placeholderArgPair.first.loc);
       BlockArgument arg = block->addArgument(placeholderArgPair.second, loc);
 
       // Add a definition of this arg to the assembly state if provided.
       if (state.asmState)
         state.asmState->addDefinition(arg, argInfo.loc);
 
       // Record the definition for this argument.
       if (addDefinition(argInfo, arg))

mlir/test/IR/locations.mlir

 }
 
 // CHECK-ALIAS: "foo.op"() : () -> () loc(#[[LOC:.*]])
 "foo.op"() : () -> () loc(#loc)
 
 // CHECK-LABEL: func @argLocs(
 // CHECK-SAME:  %arg0: i32 loc({{.*}}locations.mlir":[[# @LINE+1]]:15),
 func @argLocs(%x: i32,
-// CHECK-SAME:  %arg1: i64 loc({{.*}}locations.mlir":[[# @LINE+1]]:15))
+// CHECK-SAME:  %arg1: i64 loc("hotdog")
               %y: i64 loc("hotdog")) {
   return
 }
 
 // CHECK-LABEL: "foo.unknown_op_with_bbargs"()
 // CHECK-ALIAS: "foo.unknown_op_with_bbargs"()
 "foo.unknown_op_with_bbargs"() ({
 // CHECK-NEXT: ^bb0(%arg0: i32 loc({{.*}}locations.mlir":[[# @LINE+2]]:7),

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

static ParseResult parseIsolatedRegionOp(OpAsmParser &parser,
 
   // Parse the input operand.
   if (parser.parseOperand(argInfo) ||
       parser.resolveOperand(argInfo, argType, result.operands))
     return failure();
 
   // Parse the body region, and reuse the operand info as the argument info.
   Region *body = result.addRegion();
-  return parser.parseRegion(*body, argInfo, argType,
+  return parser.parseRegion(*body, argInfo, argType, /*argLocations=*/{},
                             /*enableNameShadowing=*/true);
 }
 
 static void print(OpAsmPrinter &p, IsolatedRegionOp op) {
   p << "test.isolated_region ";
   p.printOperand(op.getOperand());
   p.shadowRegionArgs(op.getRegion(), op.getOperand());
   p.printRegion(op.getRegion(), /*printEntryBlockArgs=*/false);