[IR] Add a Location to BlockArgument.

[lattner]

Summary

This adds the ability to specify a location when creating BlockArguments.
Notably Value::getLoc() will return this correctly, which makes diagnostics
more precise (e.g. the example in test-legalize-type-conversion.mlir).

This is currently optional to avoid breaking any existing code - if
absent, the BlockArgument defaults to using the location of its enclosing
operation (preserving existing behavior).

The bulk of this change is plumbing location tracking through the parser
and printer to make sure it can round trip (in -mlir-print-debuginfo
mode). This is complete for generic operations, but requires manual
adoption for custom ops.

I added support for function-like ops to round trip their argument
locations - they print correctly, but when parsing the locations are
dropped on the floor. I intend to fix this, but it will require more
invasive plumbing through "function_like_impl" stuff so I think it
best to split it out to its own patch.

[lattner]

[IR] Expose Block arguments through Builders.h

[rriddle]

mlir/lib/IR/Block.cpp

171–173

nit: Use llvm::zip here.

181

This comment seems a little weird. The only situation where there isn't a parent operation is if this block is unlinked.

182

nit: Spell out auto here.

mlir/lib/Parser/Parser.cpp

255

nit: Drop the llvm:: here.

[rriddle]

Do you have an estimation on memory usage growth here? Seems like a lot of existing situations would end up with unknown locations on every argument, which is a 8-byte memory size hit (per argument, in which there have been some cases in TensorFlow with hundreds of thousands of arguments) for an unused feature. For example, how often would this be used for non-entry blocks? Should we instead have a more sparse representation? e.g. a map in the parent block or something.

[rriddle]

Do you plan to resolve the TODOs in this patch or a followup?

(Code in this patch looks fine)

mlir/lib/IR/FunctionImplementation.cpp

63–67

Can we just resolve this now? Or do you plan on doing this in a followup?

[mehdi_amini]

In D102567#2764943, @rriddle wrote:

Do you have an estimation on memory usage growth here? Seems like a lot of existing situations would end up with unknown locations on every argument, which is a 8-byte memory size hit (per argument, in which there have been some cases in TensorFlow with hundreds of thousands of arguments) for an unused feature. For example, how often would this be used for non-entry blocks? Should we instead have a more sparse representation? e.g. a map in the parent block or something.

I have the same feeling here: an array of locations (lazily initialized) on the block itself would be less costly when unused.

[bondhugula]

In D102567#2764956, @mehdi_amini wrote:

In D102567#2764943, @rriddle wrote:

Do you have an estimation on memory usage growth here? Seems like a lot of existing situations would end up with unknown locations on every argument, which is a 8-byte memory size hit (per argument, in which there have been some cases in TensorFlow with hundreds of thousands of arguments) for an unused feature. For example, how often would this be used for non-entry blocks? Should we instead have a more sparse representation? e.g. a map in the parent block or something.

I have the same feeling here: an array of locations (lazily initialized) on the block itself would be less costly when unused.

Since the block argument is sort of at the same level as a defining op, having the Location on it would have been the "normal" thing. So I feel it's fine to first add this and then optimize it with indirection as suggested above later.

Side question: @rriddle did you mean "hundreds or thousands" or "hundreds of thousands" above?

[lattner]

In D102567#2764943, @rriddle wrote:

Do you have an estimation on memory usage growth here? Seems like a lot of existing situations would end up with unknown locations on every argument, which is a 8-byte memory size hit (per argument, in which there have been some cases in TensorFlow with hundreds of thousands of arguments) for an unused feature. For example, how often would this be used for non-entry blocks? Should we instead have a more sparse representation? e.g. a map in the parent block or something.

No I don't, however, I expect this to be insignificant - there are many fewer arguments than ops, attributes, and types in the primary use cases that I'm aware of. It doesn't make sense to have an argument without a use of it after all.

More generally, this patch is completely in line with MLIRs emphasis on good location tracking, and is a gaping hole in the story. We need support for this on non-entry block arguments for parity with LLVM's tracking - this is the equivalent of being able to put a location on llvm::PHINode. We also know that this causes known "error messages getting emitted to line zero" issues that Uday points out.

This is something we always planned to do, just never got around to.

[lattner]

Also, keep in mind that all block arguments coming from a .mlir file have a location. While other frontends may not be specifying them much yet, I expect that to change progressively over time.

[lattner]

mlir/lib/IR/Block.cpp

181

Good point, I just pulled it forward without rethinking it. Fixed.

mlir/lib/IR/FunctionImplementation.cpp

63–67

I plan to do it as a follow up as mentioned in the commit message above. This will require some more invasive work on function_like_impl stuff and I'd rather hill climb instead of make a single huge patch.

[lattner]

Incorporate several improvements River pointed out, thanks!

[mehdi_amini]

Since the block argument is sort of at the same level as a defining op, having the Location on it would have been the "normal" thing. So I feel it's fine to first add this and then optimize it with indirection as suggested above later.

Another way to see the analogy is that an operation defines multiple SSA values, but the operation holds a location: we don't have a location for each of the SSA values individually, and the location is owned by the entity (Operation) that encapsulate the results.

In D102567#2766214, @lattner wrote:

More generally, this patch is completely in line with MLIRs emphasis on good location tracking, and is a gaping hole in the story.

Yes we agree that this is a needed feature, but that seems orthogonal to "how do we store it".
Also the "emphasis on good location tracking" yielded APIs that forces the location to be passed in, while this makes it optional. I imagine this is just too hard to require a location for the block arguments at this point?

Also, keep in mind that all block arguments coming from a .mlir file have a location.

This does not seem like a given to me: custom parser/printer don't have to provide one for example. Would the block argument produced by the scf.for parse automatically get a location? Func op? Right now it isn't clear to me that most .mlir file would get locations on their block argument if they don't have a CFG / aren't in generic form.

No I don't, however, I expect this to be insignificant - there are many fewer arguments than ops, attributes, and types in the primary use cases that I'm aware of. It doesn't make sense to have an argument without a use of it after all.

That's convincing to me, so LGTM on this patch.

[lattner]

This does not seem like a given to me: custom parser/printer don't have to provide one for example. Would the block argument produced by the scf.for parse automatically get a location? Func op? Right now it isn't clear to me that most .mlir file would get locations on their block argument if they don't have a CFG / aren't in generic form.

Yes, fair point. I mean block arguments inside the function. My next patch in this series will fix std.func and other FunctionLike ops so we'll get good coverage for entry arguments as well.

Thank you both for the review!

[rriddle]

In D102567#2764963, @bondhugula wrote:

In D102567#2764956, @mehdi_amini wrote:

In D102567#2764943, @rriddle wrote:

Do you have an estimation on memory usage growth here? Seems like a lot of existing situations would end up with unknown locations on every argument, which is a 8-byte memory size hit (per argument, in which there have been some cases in TensorFlow with hundreds of thousands of arguments) for an unused feature. For example, how often would this be used for non-entry blocks? Should we instead have a more sparse representation? e.g. a map in the parent block or something.

I have the same feeling here: an array of locations (lazily initialized) on the block itself would be less costly when unused.

Since the block argument is sort of at the same level as a defining op, having the Location on it would have been the "normal" thing. So I feel it's fine to first add this and then optimize it with indirection as suggested above later.

Side question: @rriddle did you mean "hundreds or thousands" or "hundreds of thousands" above?

In D102567#2764963, @bondhugula wrote:

In D102567#2764956, @mehdi_amini wrote:

In D102567#2764943, @rriddle wrote:

Do you have an estimation on memory usage growth here? Seems like a lot of existing situations would end up with unknown locations on every argument, which is a 8-byte memory size hit (per argument, in which there have been some cases in TensorFlow with hundreds of thousands of arguments) for an unused feature. For example, how often would this be used for non-entry blocks? Should we instead have a more sparse representation? e.g. a map in the parent block or something.

I have the same feeling here: an array of locations (lazily initialized) on the block itself would be less costly when unused.

Since the block argument is sort of at the same level as a defining op, having the Location on it would have been the "normal" thing. So I feel it's fine to first add this and then optimize it with indirection as suggested above later.

Side question: @rriddle did you mean "hundreds or thousands" or "hundreds of thousands" above?

No, I really meant hundreds of thousands. Some are close to a million.

mlir/lib/IR/AsmPrinter.cpp

489

You need to visit the block argument locations, otherwise we don't initialize an alias for them. Take a look at void print(Operation *op) for an example. I assume we need to do that here, and when printing the block if printEntryBlockArgs is true.

[rriddle]

Noting here that I am now seeing multiple timeouts for internal TensorFlow benchmarks after this change. Compile time seems to have more than quadrupled. Looking into it now, but may request a revert.

[lattner]

Feel free to revert if this is a culprit!

I can't imagine that, but if so that's clearly the right way to go, thanks!

[rriddle]

In D102567#2766722, @lattner wrote:

Feel free to revert if this is a culprit!

I can't imagine that, but if so that's clearly the right way to go, thanks!

Thanks. One of the benchmarks has about 750k arguments, so it may just be hitting a degenerate case during parsing/printing.

[lattner]

Thanks. One of the benchmarks has about 750k arguments, so it may just be hitting a degenerate case during parsing/printing.

Whoa. Yes, I could see that exacerbating an N^2 algorithm somewhere. If you email me an mlir file that shows the issue I'd be happy to fix. Reverting is also totally fine of course.

[rriddle]

In D102567#2766826, @lattner wrote:

Thanks. One of the benchmarks has about 750k arguments, so it may just be hitting a degenerate case during parsing/printing.

Whoa. Yes, I could see that exacerbating an N^2 algorithm somewhere. If you email me an mlir file that shows the issue I'd be happy to fix. Reverting is also totally fine of course.

Can confirm after this patch just parsing the benchmark file took ~4700 seconds (1 hour and 18 minutes), with pretty much all of that being spent in Lexer::getEncodedSourceLocation. I sent out a minimal fix in D102734, which brings that time down to 8 seconds.

[lattner]

Thank you so much for working on this River!

[rriddle]

In D102567#2767482, @lattner wrote:

Thank you so much for working on this River!

I think we'll still need to revert though. Seems like there is a bug/mismatch between parsing/printing. A repro is:

func @foo(%arg0: tensor<2xi32> loc(#loc0)) {
  return
}
#loc0 = loc(unknown)

This is getting printed, but isn't able to parse back in:

/tmp/test.mlir:2:44: error: expected location instance
  func @foo(%arg0: tensor<2xi32> loc(#loc0)) {

[rriddle]

In D102567#2767539, @rriddle wrote:

In D102567#2767482, @lattner wrote:

Thank you so much for working on this River!

I think we'll still need to revert though. Seems like there is a bug/mismatch between parsing/printing. A repro is:

func @foo(%arg0: tensor<2xi32> loc(#loc0)) {
  return
}
#loc0 = loc(unknown)

This is getting printed, but isn't able to parse back in:

/tmp/test.mlir:2:44: error: expected location instance
  func @foo(%arg0: tensor<2xi32> loc(#loc0)) {

I want to say that there is some interaction going on between location aliases that are deferred (i.e. for the parent operation) and those that can't be deferred (currently for this block argument). This seems to happen when the argument inherits its parent location.

[rsmith]

In D102567#2767539, @rriddle wrote:

I think we'll still need to revert though.

Reverted for now in rG80d981eda69f1ada6d944ed89571456cad13b850.

[lattner]

Thank you for reverting this @rsmith . This is going to be very tricky to fix and may need some larger scale refactoring.

This PR isn't really required by any of my current projects, it was a "let's make MLIR better in a way that has long annoyed me" sort of thing. I'll investigate when I have time this weekend or next week. Thanks!

[rriddle]

In D102567#2768874, @lattner wrote:

Thank you for reverting this @rsmith . This is going to be very tricky to fix and may need some larger scale refactoring.

This PR isn't really required by any of my current projects, it was a "let's make MLIR better in a way that has long annoyed me" sort of thing. I'll investigate when I have time this weekend or next week. Thanks!

I can also look into the alias issue if you'd like.

[lattner]

This patch + the minimal fix for the issue identified above was reapplied in https://reviews.llvm.org/D102991. I'll continue working on improved support for this, but I'd love some help with the existing bug: https://bugs.llvm.org/show_bug.cgi?id=50451

mlir/include/mlir/IR/Block.h

 #include "mlir/IR/Visitors.h"
 
 namespace llvm {
 class BitVector;
 } // end namespace llvm
 
 namespace mlir {
 class TypeRange;
-template <typename ValueRangeT> class ValueTypeRange;
+template <typename ValueRangeT>
+class ValueTypeRange;
 
 /// `Block` represents an ordered list of `Operation`s.
 class Block : public IRObjectWithUseList<BlockOperand>,
               public llvm::ilist_node_with_parent<Block, Region> {
 public:
   explicit Block() {}
   ~Block();
 
public:
   args_iterator args_begin() { return getArguments().begin(); }
   args_iterator args_end() { return getArguments().end(); }
   reverse_args_iterator args_rbegin() { return getArguments().rbegin(); }
   reverse_args_iterator args_rend() { return getArguments().rend(); }
 
   bool args_empty() { return arguments.empty(); }
 
   /// Add one value to the argument list.
-  BlockArgument addArgument(Type type);
+  BlockArgument addArgument(Type type, Optional<Location> loc = {});
 
   /// Insert one value to the position in the argument list indicated by the
   /// given iterator. The existing arguments are shifted. The block is expected
   /// not to have predecessors.
-  BlockArgument insertArgument(args_iterator it, Type type);
+  BlockArgument insertArgument(args_iterator it, Type type,
+                               Optional<Location> loc = {});
 
   /// Add one argument to the argument list for each type specified in the list.
-  iterator_range<args_iterator> addArguments(TypeRange types);
+  iterator_range<args_iterator> addArguments(TypeRange types,
+                                             ArrayRef<Location> locs = {});
 
   /// Add one value to the argument list at the specified position.
-  BlockArgument insertArgument(unsigned index, Type type);
+  BlockArgument insertArgument(unsigned index, Type type,
+                               Optional<Location> loc = {});
 
   /// Erase the argument at 'index' and remove it from the argument list.
   void eraseArgument(unsigned index);
   /// Erases the arguments listed in `argIndices` and removes them from the
   /// argument list.
   /// `argIndices` is allowed to have duplicates and can be in any order.
   void eraseArguments(ArrayRef<unsigned> argIndices);
   /// Erases the arguments that have their corresponding bit set in
public:
 
   /// This class provides iteration over the held operations of a block for a
   /// specific operation type.
   template <typename OpT>
   using op_iterator = detail::op_iterator<OpT, iterator>;
 
   /// Return an iterator range over the operations within this block that are of
   /// 'OpT'.
-  template <typename OpT> iterator_range<op_iterator<OpT>> getOps() {
+  template <typename OpT>
+  iterator_range<op_iterator<OpT>> getOps() {
     auto endIt = end();
     return {detail::op_filter_iterator<OpT, iterator>(begin(), endIt),
             detail::op_filter_iterator<OpT, iterator>(endIt, endIt)};
   }
-  template <typename OpT> op_iterator<OpT> op_begin() {
+  template <typename OpT>
+  op_iterator<OpT> op_begin() {
     return detail::op_filter_iterator<OpT, iterator>(begin(), end());
   }
-  template <typename OpT> op_iterator<OpT> op_end() {
+  template <typename OpT>
+  op_iterator<OpT> op_end() {
     return detail::op_filter_iterator<OpT, iterator>(end(), end());
   }
 
   /// Return an iterator range over the operation within this block excluding
   /// the terminator operation at the end.
   iterator_range<iterator> without_terminator() {
     if (begin() == end())
       return {begin(), end()};

mlir/include/mlir/IR/Builders.h

public:
   IntegerType getI64Type();
   IntegerType getIntegerType(unsigned width);
   IntegerType getIntegerType(unsigned width, bool isSigned);
   FunctionType getFunctionType(TypeRange inputs, TypeRange results);
   TupleType getTupleType(TypeRange elementTypes);
   NoneType getNoneType();
 
   /// Get or construct an instance of the type 'ty' with provided arguments.
-  template <typename Ty, typename... Args> Ty getType(Args... args) {
+  template <typename Ty, typename... Args>
+  Ty getType(Args... args) {
     return Ty::get(context, args...);
   }
 
   // Attributes.
   NamedAttribute getNamedAttr(StringRef name, Attribute val);
 
   UnitAttr getUnitAttr();
   BoolAttr getBoolAttr(bool value);
public:
   //===--------------------------------------------------------------------===//
   // Block Creation
   //===--------------------------------------------------------------------===//
 
   /// Add new block with 'argTypes' arguments and set the insertion point to the
   /// end of it. The block is inserted at the provided insertion point of
   /// 'parent'.
   Block *createBlock(Region *parent, Region::iterator insertPt = {},
-                     TypeRange argTypes = llvm::None);
+                     TypeRange argTypes = llvm::None,
+                     ArrayRef<Location> locs = {});
 
   /// Add new block with 'argTypes' arguments and set the insertion point to the
   /// end of it. The block is placed before 'insertBefore'.
-  Block *createBlock(Block *insertBefore, TypeRange argTypes = llvm::None);
+  Block *createBlock(Block *insertBefore, TypeRange argTypes = llvm::None,
+                     ArrayRef<Location> locs = {});
 
   //===--------------------------------------------------------------------===//
   // Operation Creation
   //===--------------------------------------------------------------------===//
 
   /// Insert the given operation at the current insertion point and return it.
   Operation *insert(Operation *op);
 
public:
   /// empty. Operands are remapped using `mapper` (if present), and `mapper` is
   /// updated to contain the results.
   Operation *cloneWithoutRegions(Operation &op, BlockAndValueMapping &mapper) {
     return insert(op.cloneWithoutRegions(mapper));
   }
   Operation *cloneWithoutRegions(Operation &op) {
     return insert(op.cloneWithoutRegions());
   }
-  template <typename OpT> OpT cloneWithoutRegions(OpT op) {
+  template <typename OpT>
+  OpT cloneWithoutRegions(OpT op) {
     return cast<OpT>(cloneWithoutRegions(*op.getOperation()));
   }
 
 private:
   /// The current block this builder is inserting into.
   Block *block = nullptr;
   /// The insertion point within the block that this builder is inserting
   /// before.
   Block::iterator insertPoint;
   /// The optional listener for events of this builder.
   Listener *listener;
 };
 
 } // namespace mlir
 
 #endif

mlir/include/mlir/IR/OpImplementation.h

public:
   OpAsmPrinter() {}
   virtual ~OpAsmPrinter();
   virtual raw_ostream &getStream() const = 0;
 
   /// Print a newline and indent the printer to the start of the current
   /// operation.
   virtual void printNewline() = 0;
 
+  /// Print a block argument in the usual format of:
+  ///   %ssaName : type {attr1=42} loc("here")
+  /// where location printing is controlled by the standard internal option.
+  /// You may pass omitType=true to not print a type, and pass an empty
+  /// attribute list if you don't care for attributes.
+  virtual void printRegionArgument(BlockArgument arg,
+                                   ArrayRef<NamedAttribute> argAttrs = {},
+                                   bool omitType = false) = 0;
+
   /// Print implementations for various things an operation contains.
   virtual void printOperand(Value value) = 0;
   virtual void printOperand(Value value, raw_ostream &os) = 0;
 
   /// Print a comma separated list of operands.
   template <typename ContainerType>
   void printOperands(const ContainerType &container) {
     printOperands(container.begin(), container.end());
public:
   }
 
   /// Parse an optional @-identifier and store it (without the '@' symbol) in a
   /// string attribute named 'attrName'.
   virtual ParseResult parseOptionalSymbolName(StringAttr &result,
                                               StringRef attrName,
                                               NamedAttrList &attrs) = 0;
 
+  /// Parse a loc(...) specifier if present, filling in result if so.
+  virtual ParseResult
+  parseOptionalLocationSpecifier(Optional<Location> &result) = 0;
+
   //===--------------------------------------------------------------------===//
   // Operand Parsing
   //===--------------------------------------------------------------------===//
 
   /// This is the representation of an operand reference.
   struct OperandType {
     llvm::SMLoc location; // Location of the token.
     StringRef name;       // Value name, e.g. %42 or %abc

mlir/include/mlir/IR/Value.h

 /// ValueImpl that is either owner by a block(in the case of a BlockArgument) or
 /// an Operation(in the case of an OpResult).
 class Value {
 public:
   Value(detail::ValueImpl *impl = nullptr) : impl(impl) {}
   Value(const Value &) = default;
   Value &operator=(const Value &) = default;
 
-  template <typename U> bool isa() const {
+  template <typename U>
+  bool isa() const {
     assert(*this && "isa<> used on a null type.");
     return U::classof(*this);
   }
 
   template <typename First, typename Second, typename... Rest>
   bool isa() const {
     return isa<First>() || isa<Second, Rest...>();
   }
-  template <typename U> U dyn_cast() const {
+  template <typename U>
+  U dyn_cast() const {
     return isa<U>() ? U(impl) : U(nullptr);
   }
-  template <typename U> U dyn_cast_or_null() const {
+  template <typename U>
+  U dyn_cast_or_null() const {
     return (*this && isa<U>()) ? U(impl) : U(nullptr);
   }
-  template <typename U> U cast() const {
+  template <typename U>
+  U cast() const {
     assert(isa<U>());
     return U(impl);
   }
 
   explicit operator bool() const { return impl; }
   bool operator==(const Value &other) const { return impl == other.impl; }
   bool operator!=(const Value &other) const { return !(*this == other); }
 
public:
 
   /// If this value is the result of an operation of type OpTy, return the
   /// operation that defines it.
   template <typename OpTy>
   OpTy getDefiningOp() const {
     return llvm::dyn_cast_or_null<OpTy>(getDefiningOp());
   }
 
-  /// If this value is the result of an operation, use it as a location,
-  /// otherwise return an unknown location.
+  /// Return the location of this value.
   Location getLoc() const;
+  void setLoc(Location loc);
 
   /// Return the Region in which this Value is defined.
   Region *getParentRegion();
 
   /// Return the Block in which this Value is defined.
   Block *getParentBlock();
 
   //===--------------------------------------------------------------------===//

 /// The internal implementation of a BlockArgument.
 class BlockArgumentImpl : public ValueImpl {
 public:
   static bool classof(const ValueImpl *value) {
     return value->getKind() == ValueImpl::Kind::BlockArgument;
   }
 
 private:
-  BlockArgumentImpl(Type type, Block *owner, int64_t index)
-      : ValueImpl(type, Kind::BlockArgument), owner(owner), index(index) {}
+  BlockArgumentImpl(Type type, Block *owner, int64_t index, Location loc)
+      : ValueImpl(type, Kind::BlockArgument), owner(owner), index(index),
+        loc(loc) {}
 
   /// The owner of this argument.
   Block *owner;
 
   /// The position in the argument list.
   int64_t index;
 
+  /// The source location of this argument.
+  Location loc;
+
   /// Allow access to owner and constructor.
   friend BlockArgument;
 };
 } // end namespace detail
 
 /// This class represents an argument of a Block.
 class BlockArgument : public Value {
 public:
   using Value::Value;
 
   static bool classof(Value value) {
     return llvm::isa<detail::BlockArgumentImpl>(value.getImpl());
   }
 
   /// Returns the block that owns this argument.
   Block *getOwner() const { return getImpl()->owner; }
 
   /// Returns the number of this argument.
   unsigned getArgNumber() const { return getImpl()->index; }
 
+  /// Return the location for this argument.
+  Location getLoc() const { return getImpl()->loc; }
+  void setLoc(Location loc) { getImpl()->loc = loc; }
+
 private:
   /// Allocate a new argument with the given type and owner.
-  static BlockArgument create(Type type, Block *owner, int64_t index) {
-    return new detail::BlockArgumentImpl(type, owner, index);
+  static BlockArgument create(Type type, Block *owner, int64_t index,
+                              Location loc) {
+    return new detail::BlockArgumentImpl(type, owner, index, loc);
   }
 
   /// Destroy and deallocate this argument.
   void destroy() { delete getImpl(); }
 
   /// Get a raw pointer to the internal implementation.
   detail::BlockArgumentImpl *getImpl() const {
     return reinterpret_cast<detail::BlockArgumentImpl *>(impl);

 inline ::llvm::hash_code hash_value(Value arg) {
   return ::llvm::hash_value(arg.getImpl());
 }
 
 } // namespace mlir
 
 namespace llvm {
 
-template <> struct DenseMapInfo<mlir::Value> {
+template <>
+struct DenseMapInfo<mlir::Value> {
   static mlir::Value getEmptyKey() {
     void *pointer = llvm::DenseMapInfo<void *>::getEmptyKey();
     return mlir::Value::getFromOpaquePointer(pointer);
   }
   static mlir::Value getTombstoneKey() {
     void *pointer = llvm::DenseMapInfo<void *>::getTombstoneKey();
     return mlir::Value::getFromOpaquePointer(pointer);
   }
struct DenseMapInfo<mlir::BlockArgument> : public DenseMapInfo<mlir::Value> {
   }
   static mlir::BlockArgument getTombstoneKey() {
     void *pointer = llvm::DenseMapInfo<void *>::getTombstoneKey();
     return reinterpret_cast<mlir::detail::BlockArgumentImpl *>(pointer);
   }
 };
 
 /// Allow stealing the low bits of a value.
-template <> struct PointerLikeTypeTraits<mlir::Value> {
+template <>
+struct PointerLikeTypeTraits<mlir::Value> {
 public:
   static inline void *getAsVoidPointer(mlir::Value value) {
     return const_cast<void *>(value.getAsOpaquePointer());
   }
   static inline mlir::Value getFromVoidPointer(void *pointer) {
     return mlir::Value::getFromOpaquePointer(pointer);
   }
   enum {

mlir/lib/IR/AsmPrinter.cpp

void printOperand(Value, raw_ostream &os) override {
     // a value name. To maintain this invariant, emit a name even if it is
     // guaranteed to go unused.
     os << "%";
   }
   void printSymbolName(StringRef) override {}
   void printSuccessor(Block *) override {}
   void printSuccessorAndUseList(Block *, ValueRange) override {}
   void shadowRegionArgs(Region &, ValueRange) override {}
+  void printRegionArgument(BlockArgument arg, ArrayRef<NamedAttribute> argAttrs,
+                           bool omitType) override {}
 
   /// The printer flags to use when determining potential aliases.
   const OpPrintingFlags &printerFlags;
 
   /// The initializer to use when identifying aliases.
   AliasInitializer &initializer;
 
   /// A dummy output stream.
public:
   }
 
   /// Print the given attribute without its type. The corresponding parser must
   /// provide a valid type for the attribute.
   void printAttributeWithoutType(Attribute attr) override {
     ModulePrinter::printAttribute(attr, AttrTypeElision::Must);
   }
 
+  /// Print a block argument in the usual format of:
+  ///   %ssaName : type {attr1=42} loc("here")
+  /// where location printing is controlled by the standard internal option.
+  /// You may pass omitType=true to not print a type, and pass an empty
+  /// attribute list if you don't care for attributes.
+  void printRegionArgument(BlockArgument arg,
+                           ArrayRef<NamedAttribute> argAttrs = {},
+                           bool omitType = false) override;
+
   /// Print the ID for the given value.
   void printOperand(Value value) override { printValueID(value); }
   void printOperand(Value value, raw_ostream &os) override {
     printValueID(value, /*printResultNo=*/true, &os);
   }
 
   /// Print an optional attribute dictionary with a given set of elided values.
   void printOptionalAttrDict(ArrayRef<NamedAttribute> attrs,
void OperationPrinter::printTopLevelOperation(Operation *op) {
   // Print the module.
   print(op);
   os << newLine;
 
   // Output the aliases at the top level that can be deferred.
   state->getAliasState().printDeferredAliases(os, newLine);
 }
 
+/// Print a block argument in the usual format of:
+///   %ssaName : type {attr1=42} loc("here")
+/// where location printing is controlled by the standard internal option.
+/// You may pass omitType=true to not print a type, and pass an empty
+/// attribute list if you don't care for attributes.
+void OperationPrinter::printRegionArgument(BlockArgument arg,
+                                           ArrayRef<NamedAttribute> argAttrs,
+                                           bool omitType) {
+  printOperand(arg);
+  if (!omitType) {
+    os << ": ";
+    printType(arg.getType());
+  }
+  printOptionalAttrDict(argAttrs);
+  printTrailingLocation(arg.getLoc());
+}
+
 void OperationPrinter::print(Operation *op) {
   // Track the location of this operation.
   state->registerOperationLocation(op, newLine.curLine, currentIndent);
 
   os.indent(currentIndent);
   printOperation(op);
   printTrailingLocation(op->getLoc());
 }
if (printBlockArgs) {
 
     // Print the argument list if non-empty.
     if (!block->args_empty()) {
       os << '(';
       interleaveComma(block->getArguments(), [&](BlockArgument arg) {
         printValueID(arg);
         os << ": ";
         printType(arg.getType());
+        printTrailingLocation(arg.getLoc());
       });
       os << ')';
     }
     os << ':';
 
     // Print out some context information about the predecessors of this block.
     if (!block->getParent()) {
       os << "  // block is not in a region!";

 
 void IntegerSet::print(raw_ostream &os) const {
   ModulePrinter(os).printIntegerSet(*this);
 }
 
 void Value::print(raw_ostream &os) {
   if (auto *op = getDefiningOp())
     return op->print(os);
-  // TODO: Improve this.
+  // TODO: Improve BlockArgument print'ing.
   BlockArgument arg = this->cast<BlockArgument>();
   os << "<block argument> of type '" << arg.getType()
      << "' at index: " << arg.getArgNumber() << '\n';
 }
 void Value::print(raw_ostream &os, AsmState &state) {
   if (auto *op = getDefiningOp())
     return op->print(os, state);
 
-  // TODO: Improve this.
+  // TODO: Improve BlockArgument print'ing.
   BlockArgument arg = this->cast<BlockArgument>();
   os << "<block argument> of type '" << arg.getType()
      << "' at index: " << arg.getArgNumber() << '\n';
 }
 
 void Value::dump() {
   print(llvm::errs());
   llvm::errs() << "\n";

mlir/lib/IR/Block.cpp

 // Argument list management.
 //===----------------------------------------------------------------------===//
 
 /// Return a range containing the types of the arguments for this block.
 auto Block::getArgumentTypes() -> ValueTypeRange<BlockArgListType> {
   return ValueTypeRange<BlockArgListType>(getArguments());
 }
 
-BlockArgument Block::addArgument(Type type) {
-  BlockArgument arg = BlockArgument::create(type, this, arguments.size());
+BlockArgument Block::addArgument(Type type, Optional<Location> loc) {
+  // TODO: Require locations for BlockArguments.
+  if (!loc.hasValue()) {
+    // Use the location of the parent operation if the block is attached.
+    if (Operation *parentOp = getParentOp())
+      loc = parentOp->getLoc();
+    else
+      loc = UnknownLoc::get(type.getContext());
+  }
+
+  BlockArgument arg = BlockArgument::create(type, this, arguments.size(), *loc);
   arguments.push_back(arg);
   return arg;
 }
 
 /// Add one argument to the argument list for each type specified in the list.
-auto Block::addArguments(TypeRange types) -> iterator_range<args_iterator> {
+auto Block::addArguments(TypeRange types, ArrayRef<Location> locs)
+    -> iterator_range<args_iterator> {
+  // TODO: Require locations for BlockArguments.
+  assert((locs.empty() || types.size() == locs.size()) &&
+         "incorrect number of block argument locations");
   size_t initialSize = arguments.size();
+
   arguments.reserve(initialSize + types.size());
+
+  // TODO: Require locations for BlockArguments.
+  if (locs.empty()) {
     for (auto type : types)
       addArgument(type);
+  } else {
+    for (auto typeAndLoc : llvm::zip(types, locs))
+      addArgument(std::get<0>(typeAndLoc), std::get<1>(typeAndLoc));
+  }
   return {arguments.data() + initialSize, arguments.data() + arguments.size()};
 }
 
-BlockArgument Block::insertArgument(unsigned index, Type type) {
-  auto arg = BlockArgument::create(type, this, index);
+BlockArgument Block::insertArgument(unsigned index, Type type,
+                                    Optional<Location> loc) {
+  // TODO: Require locations for BlockArguments.
+  if (!loc.hasValue()) {
+    // Use the location of the parent operation if the block is attached.
+    if (Operation *parentOp = getParentOp())
+      loc = parentOp->getLoc();
+    else
+      loc = UnknownLoc::get(type.getContext());
+  }
+
+  auto arg = BlockArgument::create(type, this, index, *loc);
   assert(index <= arguments.size());
   arguments.insert(arguments.begin() + index, arg);
   // Update the cached position for all the arguments after the newly inserted
   // one.
   ++index;
   for (BlockArgument arg : llvm::drop_begin(arguments, index))
     arg.setArgNumber(index++);
   return arg;
 }
 
 /// Insert one value to the given position of the argument list. The existing
 /// arguments are shifted. The block is expected not to have predecessors.
-BlockArgument Block::insertArgument(args_iterator it, Type type) {
+BlockArgument Block::insertArgument(args_iterator it, Type type,
+                                    Optional<Location> loc) {
   assert(llvm::empty(getPredecessors()) &&
          "cannot insert arguments to blocks with predecessors");
-  return insertArgument(it->getArgNumber(), type);
+  return insertArgument(it->getArgNumber(), type, loc);
 }
 
 void Block::eraseArgument(unsigned index) {
   assert(index < arguments.size());
   arguments[index].destroy();
   arguments.erase(arguments.begin() + index);
   for (BlockArgument arg : llvm::drop_begin(arguments, index))
     arg.setArgNumber(index++);

mlir/lib/IR/Builders.cpp

if (listener)
     listener->notifyOperationInserted(op);
   return op;
 }
 
 /// Add new block with 'argTypes' arguments and set the insertion point to the
 /// end of it. The block is inserted at the provided insertion point of
 /// 'parent'.
 Block *OpBuilder::createBlock(Region *parent, Region::iterator insertPt,
-                              TypeRange argTypes) {
+                              TypeRange argTypes, ArrayRef<Location> locs) {
   assert(parent && "expected valid parent region");
   if (insertPt == Region::iterator())
     insertPt = parent->end();
 
   Block *b = new Block();
-  b->addArguments(argTypes);
+  b->addArguments(argTypes, locs);
   parent->getBlocks().insert(insertPt, b);
   setInsertionPointToEnd(b);
 
   if (listener)
     listener->notifyBlockCreated(b);
   return b;
 }
 
 /// Add new block with 'argTypes' arguments and set the insertion point to the
 /// end of it.  The block is placed before 'insertBefore'.
-Block *OpBuilder::createBlock(Block *insertBefore, TypeRange argTypes) {
+Block *OpBuilder::createBlock(Block *insertBefore, TypeRange argTypes,
+                              ArrayRef<Location> locs) {
   assert(insertBefore && "expected valid insertion block");
   return createBlock(insertBefore->getParent(), Region::iterator(insertBefore),
-                     argTypes);
+                     argTypes, locs);
 }
 
 /// Create an operation given the fields represented as an OperationState.
 Operation *OpBuilder::createOperation(const OperationState &state) {
   return insert(Operation::create(state));
 }
 
 /// Attempts to fold the given operation and places new results within

mlir/lib/IR/FunctionImplementation.cpp

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.
+    Optional<Location> explicitLoc;
+    if (!argument.name.empty() &&
+        parser.parseOptionalLocationSpecifier(explicitLoc))
+      return failure();
+
     return success();
   };
 
   // Parse the function arguments.
   isVariadic = false;
   if (failed(parser.parseOptionalRParen())) {
     do {
       unsigned numTypedArguments = argTypes.size();
void mlir::function_like_impl::printFunctionSignature(
 
   p << '(';
   ArrayAttr argAttrs = op->getAttrOfType<ArrayAttr>(getArgDictAttrName());
   for (unsigned i = 0, e = argTypes.size(); i < e; ++i) {
     if (i > 0)
       p << ", ";
 
     if (!isExternal) {
-      p.printOperand(body.getArgument(i));
-      p << ": ";
-    }
-
+      ArrayRef<NamedAttribute> attrs;
+      if (argAttrs)
+        attrs = argAttrs[i].cast<DictionaryAttr>().getValue();
+      p.printRegionArgument(body.getArgument(i), attrs);
+    } else {
       p.printType(argTypes[i]);
       if (argAttrs)
         p.printOptionalAttrDict(argAttrs[i].cast<DictionaryAttr>().getValue());
     }
+  }
 
   if (isVariadic) {
     if (!argTypes.empty())
       p << ", ";
     p << "...";
   }
 
   p << ')';

mlir/lib/IR/Value.cpp

if (auto result = dyn_cast<OpResult>())
     return result.getOwner();
   return nullptr;
 }
 
 Location Value::getLoc() const {
   if (auto *op = getDefiningOp())
     return op->getLoc();
 
-  // Use the location of the parent operation if this is a block argument.
-  // TODO: Should we just add locations to block arguments?
-  Operation *parentOp = cast<BlockArgument>().getOwner()->getParentOp();
-  return parentOp ? parentOp->getLoc() : UnknownLoc::get(getContext());
+  return cast<BlockArgument>().getLoc();
+}
+
+void Value::setLoc(Location loc) {
+  if (auto *op = getDefiningOp())
+    return op->setLoc(loc);
+
+  return cast<BlockArgument>().setLoc(loc);
 }
 
 /// Return the Region in which this Value is defined.
 Region *Value::getParentRegion() {
   if (auto *op = getDefiningOp())
     return op->getParentRegion();
   return cast<BlockArgument>().getOwner()->getParent();
 }

mlir/lib/Parser/Parser.cpp

public:
   /// Parse an operation instance that is in the generic form.
   Operation *parseGenericOperation();
 
   /// Parse an operation instance that is in the generic form and insert it at
   /// the provided insertion point.
   Operation *parseGenericOperation(Block *insertBlock,
                                    Block::iterator insertPt);
 
-  /// Parse an optional trailing location for the given operation.
+  /// This type is used to keep track of things that are either an Operation or
+  /// a BlockArgument.  We cannot use Value for this, because not all Operations
+  /// have results.
+  using OpOrArgument = llvm::PointerUnion<Operation *, BlockArgument>;
+
+  /// Parse an optional trailing location and add it to the specifier Operation
+  /// or `OperandType` if present.
   ///
   ///   trailing-location ::= (`loc` (`(` location `)` | attribute-alias))?
   ///
-  ParseResult parseTrailingOperationLocation(Operation *op);
+  ParseResult parseTrailingLocationSpecifier(OpOrArgument opOrArgument);
 
   /// This is the structure of a result specifier in the assembly syntax,
   /// including the name, number of results, and location.
   using ResultRecord = std::tuple<StringRef, unsigned, SMLoc>;
 
   /// Parse an operation instance that is in the op-defined custom form.
   /// resultInfo specifies information about the "%name =" specifiers.
   Operation *parseCustomOperation(ArrayRef<ResultRecord> resultIDs);
private:
   SmallVector<DenseMap<Block *, SMLoc>, 2> forwardRef;
 
   /// These are all of the placeholders we've made along with the location of
   /// their first reference, to allow checking for use of undefined values.
   DenseMap<Value, SMLoc> forwardRefPlaceholders;
 
   /// A set of operations whose locations reference aliases that have yet to
   /// be resolved.
-  SmallVector<std::pair<Operation *, Token>, 8> opsWithDeferredLocs;
+  SmallVector<std::pair<OpOrArgument, Token>, 8>
+      opsAndArgumentsWithDeferredLocs;
 
   /// The builder used when creating parsed operation instances.
   OpBuilder opBuilder;
 
   /// The top level operation that holds all of the parsed operations.
   Operation *topLevelOp;
 };
 } // end anonymous namespace
for (auto entry : errors) {
       auto loc = SMLoc::getFromPointer(entry);
       emitError(loc, "use of undeclared SSA value name");
     }
     return failure();
   }
 
   // Resolve the locations of any deferred operations.
   auto &attributeAliases = state.symbols.attributeAliasDefinitions;
-  for (std::pair<Operation *, Token> &it : opsWithDeferredLocs) {
+  for (std::pair<OpOrArgument, Token> &it : opsAndArgumentsWithDeferredLocs) {
     llvm::SMLoc tokLoc = it.second.getLoc();
     StringRef identifier = it.second.getSpelling().drop_front();
     Attribute attr = attributeAliases.lookup(identifier);
     if (!attr)
       return emitError(tokLoc) << "operation location alias was never defined";
 
     LocationAttr locAttr = attr.dyn_cast<LocationAttr>();
     if (!locAttr)
       return emitError(tokLoc)
              << "expected location, but found '" << attr << "'";
-    it.first->setLoc(locAttr);
+    auto opOrArgument = it.first;
+    if (auto *op = opOrArgument.dyn_cast<Operation *>())
+      op->setLoc(locAttr);
+    else
+      opOrArgument.get<BlockArgument>().setLoc(locAttr);
   }
 
   // Pop the top level name scope.
   return popSSANameScope();
 }
 
 //===----------------------------------------------------------------------===//
 // SSA Value Handling
Operation *OperationParser::parseGenericOperation() {
   for (unsigned i = 0, e = operandInfos.size(); i != e; ++i) {
     result.operands.push_back(resolveSSAUse(operandInfos[i], operandTypes[i]));
     if (!result.operands.back())
       return nullptr;
   }
 
   // Create the operation and try to parse a location for it.
   Operation *op = opBuilder.createOperation(result);
-  if (parseTrailingOperationLocation(op))
+  if (parseTrailingLocationSpecifier(op))
     return nullptr;
   return op;
 }
 
 Operation *OperationParser::parseGenericOperation(Block *insertBlock,
                                                   Block::iterator insertPt) {
   OpBuilder::InsertionGuard restoreInsertionPoint(opBuilder);
   opBuilder.setInsertionPoint(insertBlock, insertPt);
if (atToken.isNot(Token::at_identifier))
       return failure();
 
     result = getBuilder().getStringAttr(atToken.getSymbolReference());
     attrs.push_back(getBuilder().getNamedAttr(attrName, result));
     parser.consumeToken();
     return success();
   }
 
+  /// Parse a loc(...) specifier if present, filling in result if so.
+  ParseResult
+  parseOptionalLocationSpecifier(Optional<Location> &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();
+  }
+
   //===--------------------------------------------------------------------===//
   // Operand Parsing
   //===--------------------------------------------------------------------===//
 
   /// Parse a single operand.
   ParseResult parseOperand(OperandType &result) override {
     OperationParser::SSAUseInfo useInfo;
     if (parser.parseSSAUse(useInfo))
if (opAsmParser.parseOperation(opState))
     return nullptr;
 
   // If it emitted an error, we failed.
   if (opAsmParser.didEmitError())
     return nullptr;
 
   // Otherwise, create the operation and try to parse a location for it.
   Operation *op = opBuilder.createOperation(opState);
-  if (parseTrailingOperationLocation(op))
+  if (parseTrailingLocationSpecifier(op))
     return nullptr;
   return op;
 }
 
-ParseResult OperationParser::parseTrailingOperationLocation(Operation *op) {
+ParseResult
+OperationParser::parseTrailingLocationSpecifier(OpOrArgument opOrArgument) {
   // If there is a 'loc' we parse a trailing location.
   if (!consumeIf(Token::kw_loc))
     return success();
   if (parseToken(Token::l_paren, "expected '(' in location"))
     return failure();
   Token tok = getToken();
 
   // Check to see if we are parsing a location alias.
if (tok.is(Token::hash_identifier)) {
     // If this alias can be resolved, do it now.
     Attribute attr = state.symbols.attributeAliasDefinitions.lookup(identifier);
     if (attr) {
       if (!(directLoc = attr.dyn_cast<LocationAttr>()))
         return emitError(tok.getLoc())
                << "expected location, but found '" << attr << "'";
     } else {
       // Otherwise, remember this operation and resolve its location later.
-      opsWithDeferredLocs.emplace_back(op, tok);
+      opsAndArgumentsWithDeferredLocs.emplace_back(opOrArgument, tok);
     }
 
     // Otherwise, we parse the location directly.
   } else if (parseLocationInstance(directLoc)) {
     return failure();
   }
 
   if (parseToken(Token::r_paren, "expected ')' in location"))
     return failure();
 
-  if (directLoc)
+  if (directLoc) {
+    if (auto *op = opOrArgument.dyn_cast<Operation *>())
       op->setLoc(directLoc);
+    else
+      opOrArgument.get<BlockArgument>().setLoc(directLoc);
+  }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Region Parsing
 //===----------------------------------------------------------------------===//
 
 /// Region.
for (auto &placeholderArgPair : entryArguments) {
 
       // 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";
       }
-      BlockArgument arg = block->addArgument(placeholderArgPair.second);
+      auto loc = 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))
         return failure();
return parseSSADefOrUseAndType(
             if (nextArgument >= owner->getNumArguments())
               return emitError("too many arguments specified in argument list");
 
             // Finally, make sure the existing argument has the correct type.
             arg = owner->getArgument(nextArgument++);
             if (arg.getType() != type)
               return emitError("argument and block argument type mismatch");
           } else {
-            arg = owner->addArgument(type);
+            auto loc = getEncodedSourceLocation(useInfo.loc);
+            arg = owner->addArgument(type, loc);
           }
 
+          // If the argument has an explicit loc(...) specifier, parse and apply
+          // it.
+          if (parseTrailingLocationSpecifier(arg))
+            return failure();
+
           // Mark this block argument definition in the parser state if it was
           // provided.
           if (state.asmState)
             state.asmState->addDefinition(arg, useInfo.loc);
 
           return addDefinition(useInfo, arg);
         });
   });

mlir/test/IR/locations.mlir

func @escape_strings() {
   // CHECK: loc("escaped\0A":0:0)
   "foo"() : () -> () loc("escaped\n":0:0)
   return
 }
 
 // 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))
+              %y: i64 loc("hotdog")) {
+  return
+}
+
+// CHECK-LABEL: "foo.unknown_op_with_bbargs"()
+"foo.unknown_op_with_bbargs"() ({
+// CHECK-NEXT: ^bb0(%arg0: i32 loc({{.*}}locations.mlir":[[# @LINE+1]]:7),
+ ^bb0(%x: i32,
+// CHECK-SAME: %arg1: i32 loc("cheetos"),
+      %y: i32 loc("cheetos"),
+// CHECK-SAME: %arg2: i32 loc("out_of_line_location")):
+      %z: i32 loc(#loc)):
+    %1 = addi %x, %y : i32
+    "foo.yield"(%1) : (i32) -> ()
+  }) : () -> ()
+
 // CHECK-ALIAS: #[[LOC]] = loc("out_of_line_location")
 #loc = loc("out_of_line_location")

mlir/test/Transforms/test-legalize-type-conversion.mlir

 // RUN: mlir-opt %s -test-legalize-type-conversion -allow-unregistered-dialect -split-input-file -verify-diagnostics | FileCheck %s
 
+
+func @test_invalid_arg_materialization(
   // expected-error@below {{failed to materialize conversion for block argument #0 that remained live after conversion, type was 'i16'}}
-func @test_invalid_arg_materialization(%arg0: i16) {
+  %arg0: i16) {
   // expected-note@below {{see existing live user here}}
   "foo.return"(%arg0) : (i16) -> ()
 }
 
 // -----
 
 // expected-error@below {{failed to legalize conversion operation generated for block argument}}
 func @test_invalid_arg_illegal_materialization(%arg0: i32) {

mlir/test/mlir-tblgen/pattern.mlir

 // CHECK-LABEL: verifyZeroArg
 func @verifyZeroArg() -> i32 {
   // CHECK: "test.op_k"() : () -> i32
   %0 = "test.op_j"() : () -> i32
   return %0 : i32
 }
 
 // CHECK-LABEL: testIgnoreArgMatch
-// CHECK-SAME: (%{{[a-z0-9]*}}: i32, %[[ARG1:[a-z0-9]*]]: i32
+// CHECK-SAME: (%{{[a-z0-9]*}}: i32 loc({{[^)]*}}), %[[ARG1:[a-z0-9]*]]: i32 loc({{[^)]*}}),
 func @testIgnoreArgMatch(%arg0: i32, %arg1: i32, %arg2: i32, %arg3: f32) {
   // CHECK: "test.ignore_arg_match_dst"(%[[ARG1]]) {f = 15 : i64}
   "test.ignore_arg_match_src"(%arg0, %arg1, %arg2) {d = 42, e = 24, f = 15} : (i32, i32, i32) -> ()
 
   // CHECK: test.ignore_arg_match_src
   // Not match because wrong type for $c.
   "test.ignore_arg_match_src"(%arg0, %arg1, %arg3) {d = 42, e = 24, f = 15} : (i32, i32, f32) -> ()
 
   // CHECK: test.ignore_arg_match_src
   // Not match because wrong type for $f.
   "test.ignore_arg_match_src"(%arg0, %arg1, %arg2) {d = 42 : i32, e = 24, f = 15} : (i32, i32, i32) -> ()
   return
 }
 
 // CHECK-LABEL: verifyInterleavedOperandAttribute
-// CHECK-SAME:    %[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32
+// CHECK-SAME:    %[[ARG0:.*]]: i32 loc({{[^)]*}}), %[[ARG1:.*]]: i32 loc({{[^)]*}})
 func @verifyInterleavedOperandAttribute(%arg0: i32, %arg1: i32) {
   // CHECK: "test.interleaved_operand_attr2"(%[[ARG0]], %[[ARG1]]) {attr1 = 15 : i64, attr2 = 42 : i64}
   "test.interleaved_operand_attr1"(%arg0, %arg1) {attr1 = 15, attr2 = 42} : (i32, i32) -> ()
   return
 }
 
 // CHECK-LABEL: verifyBenefit
 func @verifyBenefit(%arg0 : i32) -> i32 {
func @verifyAllAttrConstraintOf() -> (i32, i32, i32) {
   // CHECK: "test.all_attr_constraint_of1"
   %1 = "test.all_attr_constraint_of1"() {attr = [0, 2]} : () -> (i32)
   // CHECK: "test.all_attr_constraint_of1"
   %2 = "test.all_attr_constraint_of1"() {attr = [-1, 1]} : () -> (i32)
   return %0, %1, %2: i32, i32, i32
 }
 
 // CHECK-LABEL: verifyManyArgs
-// CHECK-SAME: (%[[ARG:.*]]: i32)
+// CHECK-SAME: (%[[ARG:.*]]: i32 loc({{[^)]*}}))
 func @verifyManyArgs(%arg: i32) {
   // CHECK: "test.many_arguments"(%[[ARG]], %[[ARG]], %[[ARG]], %[[ARG]], %[[ARG]], %[[ARG]], %[[ARG]], %[[ARG]], %[[ARG]])
   // CHECK-SAME: {attr1 = 24 : i64, attr2 = 42 : i64, attr3 = 42 : i64, attr4 = 42 : i64, attr5 = 42 : i64, attr6 = 42 : i64, attr7 = 42 : i64, attr8 = 42 : i64, attr9 = 42 : i64}
   "test.many_arguments"(%arg, %arg, %arg, %arg, %arg, %arg, %arg, %arg, %arg) {
     attr1 = 42, attr2 = 42, attr3 = 42, attr4 = 42, attr5 = 42,
     attr6 = 42, attr7 = 42, attr8 = 42, attr9 = 42
   } : (i32, i32, i32, i32, i32, i32, i32, i32, i32) -> ()
   return