[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

diff --git a/mlir/include/mlir/IR/FunctionImplementation.h b/mlir/include/mlir/IR/FunctionImplementation.h
index 3d65452..04684e8 100644 (file)
--- a/mlir/include/mlir/IR/FunctionImplementation.h
+++ b/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,
-    SmallVectorImpl<Type> &resultTypes,
+    bool &isVariadic, SmallVectorImpl<Type> &resultTypes,
     SmallVectorImpl<NamedAttrList> &resultAttrs);
 
 /// Parser implementation for function-like operations.  Uses

diff --git a/mlir/include/mlir/IR/OpImplementation.h b/mlir/include/mlir/IR/OpImplementation.h
index 2cdce8b..2d0a5a8 100644 (file)
--- a/mlir/include/mlir/IR/OpImplementation.h
+++ b/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
-  /// non-empty it contains a 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'.
+  /// 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<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 = {},
-      bool enableNameShadowing = false) = 0;
+      ArrayRef<Type> argTypes = {}, 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,

diff --git a/mlir/lib/Dialect/Async/IR/Async.cpp b/mlir/lib/Dialect/Async/IR/Async.cpp
index 55f7f7a..fb64c39 100644 (file)
--- a/mlir/lib/Dialect/Async/IR/Async.cpp
+++ b/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();
 

diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index 9de7c7e..255ae0e 100644 (file)
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/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,
-      isVariadic);
+      /*allowVariadic=*/false, argNames, argTypes, argAttrs, 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())

diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index 4211828..a20b637 100644 (file)
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/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 =

diff --git a/mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp b/mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp
index 14d7f78..337bd41 100644 (file)
--- a/mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp
+++ b/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);

diff --git a/mlir/lib/IR/FunctionImplementation.cpp b/mlir/lib/IR/FunctionImplementation.cpp
index 2725d48..e33150d 100644 (file)
--- a/mlir/lib/IR/FunctionImplementation.cpp
+++ b/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.
-    Optional<Location> explicitLoc;
-    if (!argument.name.empty() &&
-        parser.parseOptionalLocationSpecifier(explicitLoc))
-      return failure();
-    if (!explicitLoc)
-      explicitLoc = parser.getEncodedSourceLoc(loc);
-    argLocations.push_back(*explicitLoc);
-
+    // If we had an argument name, then remember the parsed argument.
+    if (!argument.name.empty())
+      argNames.push_back(argument);
     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,
-    SmallVectorImpl<Type> &resultTypes,
+    bool &isVariadic, 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,
-                             resultAttrs))
+                             argAttrs, isVariadic, resultTypes, 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))

diff --git a/mlir/lib/Parser/Parser.cpp b/mlir/lib/Parser/Parser.cpp
index 22daae9..9aa6976 100644 (file)
--- a/mlir/lib/Parser/Parser.cpp
+++ b/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
-  /// argument. 'isIsolatedNameScope' indicates if the naming scope of this
-  /// region is isolated from those above.
+  /// 'isIsolatedNameScope' indicates if the naming scope of this 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=*/{},
-                        /*argLocations=*/{}))
+        if (parseRegion(*result.regions.back(), /*entryArguments=*/{}))
           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};
-    return success();
+    result = {useInfo.location, useInfo.name, useInfo.number, {}};
+
+    // 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,
-                           enableNameShadowing))
+    if (parser.parseRegion(region, regionArguments, 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,
-                       enableNameShadowing);
+    return parseRegion(region, arguments, argTypes, 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=*/{},
-                    enableNameShadowing))
+    if (parseRegion(*newRegion, arguments, argTypes, 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) {
-  assert(argLocations.empty() || argLocations.size() == entryArguments.size());
+    bool isIsolatedNameScope) {
   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)) {
-      size_t argIndex = it.index();
-      auto &placeholderArgPair = it.value();
+    for (auto &placeholderArgPair : entryArguments) {
       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(
-          placeholderArgPair.second,
-          argLocations.empty()
-              ? getEncodedSourceLocation(placeholderArgPair.first.location)
-              : argLocations[argIndex]);
+      Location loc = argInfo.sourceLoc.hasValue()
+                         ? argInfo.sourceLoc.getValue()
+                         : getEncodedSourceLocation(argInfo.location);
+      BlockArgument arg = block->addArgument(placeholderArgPair.second, loc);
 
       // Add a definition of this arg to the assembly state if provided.
       if (state.asmState)

diff --git a/mlir/test/Dialect/LLVMIR/func.mlir b/mlir/test/Dialect/LLVMIR/func.mlir
index d8761e4..810622d 100644 (file)
--- a/mlir/test/Dialect/LLVMIR/func.mlir
+++ b/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
   }
 

diff --git a/mlir/test/lib/Dialect/Test/TestDialect.cpp b/mlir/test/lib/Dialect/Test/TestDialect.cpp
index 5cfefbb..c7aa3f1 100644 (file)
--- a/mlir/test/lib/Dialect/Test/TestDialect.cpp
+++ b/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);
 }