[mlir] Add support for generating debug locations from intermediate levels of the IR.

[rriddle]

Summary

This revision adds a utility to generate debug locations from the IR during compilation, by snapshotting to a output stream and using the locations that operations were dumped in that stream. The new locations may either;

• Replace the original location of the operation.

old:

loc("original_source.cpp":1:1)

new:

loc("snapshot_source.mlir":10:10)

• Fuse with the original locations as NamedLocs with a specific tag.

old:

loc("original_source.cpp":1:1)

new:

loc(fused["original_source.cpp":1:1, "snapshot"("snapshot_source.mlir":10:10)])

This feature may be used by a debugger to display the code at various different levels of the IR. It would also be able to show the different levels of IR attached to a specific source line in the original source file.

This feature may also be used to generate locations for operations generated during compilation, that don't necessarily have a user source location to attach to.

This requires changes in the printer to track the locations of operations emitted in the stream. Moving forward we need to properly(and efficiently) track the number of newlines emitted to the stream during printing.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: unknown.

Build artifacts: diff.json, CMakeCache.txt, console-log.txt

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[mehdi_amini]

Nice!

mlir/lib/Transforms/LocationSnapshot.cpp

31

Can this all move after the printing?

43

Is this expected / possible? Otherwise assert()? (if this is a valid case, can you add a comment?)

[rriddle]

Resolve comments

mlir/lib/Transforms/LocationSnapshot.cpp

43

Yep, some region terminators are elided so they won't have a valid custom assembly form. Added a comment.

mlir/include/mlir/IR/AsmState.h

 //===- AsmState.h - State class for AsmPrinter ------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines the AsmState class.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_ASMSTATE_H_
 #define MLIR_IR_ASMSTATE_H_
 
+#include "mlir/Support/LLVM.h"
+
 #include <memory>
 
 namespace mlir {
 class Operation;
 
 namespace detail {
 class AsmStateImpl;
 } // end namespace detail
 
 /// This class provides management for the lifetime of the state used when
 /// printing the IR. It allows for alleviating the cost of recomputing the
 /// internal state of the asm printer.
 ///
 /// The IR should not be mutated in-between invocations using this state, and
 /// the IR being printed must not be an parent of the IR originally used to
 /// initialize this state. This means that if a child operation is provided, a
 /// parent operation cannot reuse this state.
 class AsmState {
 public:
-  /// Initialize the asm state at the level of the given operation.
-  AsmState(Operation *op);
+  /// This map represents the raw locations of operations within the output
+  /// stream. This maps the original pointer to the operation, to a pair of line
+  /// and column in the output stream.
+  using LocationMap = DenseMap<Operation *, std::pair<unsigned, unsigned>>;
+
+  /// Initialize the asm state at the level of the given operation. A location
+  /// map may optionally be provided to be populated when printing.
+  AsmState(Operation *op, LocationMap *locationMap = nullptr);
   ~AsmState();
 
   /// Return an instance of the internal implementation. Returns nullptr if the
   /// state has not been initialized.
   detail::AsmStateImpl &getImpl() { return *impl; }
 
 private:
   AsmState() = delete;
 
   /// A pointer to allocated storage for the impl state.
   std::unique_ptr<detail::AsmStateImpl> impl;
 };
 
 } // end namespace mlir
 
 #endif // MLIR_IR_ASMSTATE_H_

mlir/include/mlir/Transforms/LocationSnapshot.h

+//===- LocationSnapshot.h - Location Snapshot Utilities ---------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This header file several utility methods for snapshotting the current IR to
+// produce new debug locations.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TRANSFORMS_LOCATIONSNAPSHOT_H
+#define MLIR_TRANSFORMS_LOCATIONSNAPSHOT_H
+
+#include "mlir/Support/LLVM.h"
+#include "llvm/ADT/StringRef.h"
+
+#include <memory>
+
+namespace mlir {
+class Location;
+struct LogicalResult;
+class Operation;
+class OpPrintingFlags;
+class Pass;
+
+/// This function generates new locations from the given IR by snapshotting the
+/// IR to the given stream, and using the printed locations within that stream.
+/// The generated locations replace the current operation locations.
+void generateLocationsFromIR(raw_ostream &os, StringRef fileName, Operation *op,
+                             OpPrintingFlags flags);
+/// This function generates new locations from the given IR by snapshotting the
+/// IR to the given file, and using the printed locations within that file. If
+/// `filename` is empty, a temporary file is generated instead.
+LogicalResult generateLocationsFromIR(StringRef fileName, Operation *op,
+                                      OpPrintingFlags flags);
+
+/// This function generates new locations from the given IR by snapshotting the
+/// IR to the given stream, and using the printed locations within that stream.
+/// The generated locations are represented as a NameLoc with the given tag as
+/// the name, and then fused with the existing locations.
+void generateLocationsFromIR(raw_ostream &os, StringRef fileName, StringRef tag,
+                             Operation *op, OpPrintingFlags flags);
+/// This function generates new locations from the given IR by snapshotting the
+/// IR to the given file, and using the printed locations within that file. If
+/// `filename` is empty, a temporary file is generated instead.
+LogicalResult generateLocationsFromIR(StringRef fileName, StringRef tag,
+                                      Operation *op, OpPrintingFlags flags);
+
+/// Create a pass to generate new locations by snapshotting the IR to the given
+/// file, and using the printed locations within that file. If `filename` is
+/// empty, a temporary file is generated instead. If a 'tag' is non-empty, the
+/// generated locations are represented as a NameLoc with the given tag as the
+/// name, and then fused with the existing locations. Otherwise, the existing
+/// locations are replaced.
+std::unique_ptr<Pass> createLocationSnapshotPass(OpPrintingFlags flags,
+                                                 StringRef fileName = "",
+                                                 StringRef tag = "");
+
+} // end namespace mlir
+
+#endif // MLIR_TRANSFORMS_LOCATIONSNAPSHOT_H

mlir/lib/IR/AsmPrinter.cpp

 bool OpPrintingFlags::shouldPrintGenericOpForm() const {
   return printGenericOpFormFlag;
 }
 
 /// Return if the printer should use local scope when dumping the IR.
 bool OpPrintingFlags::shouldUseLocalScope() const { return printLocalScope; }
 
 //===----------------------------------------------------------------------===//
+// NewLineCounter
+//===----------------------------------------------------------------------===//
+
+namespace {
+/// This class is a simple formatter that emits a new line when inputted into a
+/// stream, that enables counting the number of newlines emitted. This class
+/// should be used whenever emitting newlines in the printer.
+struct NewLineCounter {
+  unsigned curLine = 1;
+};
+} // end anonymous namespace
+
+static raw_ostream &operator<<(raw_ostream &os, NewLineCounter &newLine) {
+  ++newLine.curLine;
+  return os << '\n';
+}
+
+//===----------------------------------------------------------------------===//
 // AliasState
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class manages the state for type and attribute aliases.
 class AliasState {
 public:
   // Initialize the internal aliases.
   void
   initialize(Operation *op,
              DialectInterfaceCollection<OpAsmDialectInterface> &interfaces);
 
   /// Return a name used for an attribute alias, or empty if there is no alias.
   Twine getAttributeAlias(Attribute attr) const;
 
   /// Print all of the referenced attribute aliases.
-  void printAttributeAliases(raw_ostream &os) const;
+  void printAttributeAliases(raw_ostream &os, NewLineCounter &newLine) const;
 
   /// Return a string to use as an alias for the given type, or empty if there
   /// is no alias recorded.
   StringRef getTypeAlias(Type ty) const;
 
   /// Print all of the referenced type aliases.
-  void printTypeAliases(raw_ostream &os) const;
+  void printTypeAliases(raw_ostream &os, NewLineCounter &newLine) const;
 
 private:
   /// A special index constant used for non-kind attribute aliases.
   enum { NonAttrKindAlias = -1 };
 
   /// Record a reference to the given attribute.
   void recordAttributeReference(Attribute attr);
 
Twine AliasState::getAttributeAlias(Attribute attr) const {
   // Return the alias for this attribute, along with the index if this was
   // generated by a kind alias.
   int kindIndex = alias->second.second;
   return alias->second.first +
          (kindIndex == NonAttrKindAlias ? Twine() : Twine(kindIndex));
 }
 
 /// Print all of the referenced attribute aliases.
-void AliasState::printAttributeAliases(raw_ostream &os) const {
+void AliasState::printAttributeAliases(raw_ostream &os,
+                                       NewLineCounter &newLine) const {
   auto printAlias = [&](StringRef alias, Attribute attr, int index) {
     os << '#' << alias;
     if (index != NonAttrKindAlias)
       os << index;
-    os << " = " << attr << '\n';
+    os << " = " << attr << newLine;
   };
 
   // Print all of the attribute kind aliases.
   for (auto &kindAlias : attrKindToAlias) {
     auto &aliasAttrsPair = kindAlias.second;
     for (unsigned i = 0, e = aliasAttrsPair.second.size(); i != e; ++i)
       printAlias(aliasAttrsPair.first, aliasAttrsPair.second[i], i);
-    os << "\n";
+    os << newLine;
   }
 
   // In a second pass print all of the remaining attribute aliases that aren't
   // kind aliases.
   for (Attribute attr : usedAttributes) {
     auto alias = attrToAlias.find(attr);
     if (alias != attrToAlias.end() && alias->second.second == NonAttrKindAlias)
       printAlias(alias->second.first, attr, alias->second.second);
   }
 }
 
 /// Return a string to use as an alias for the given type, or empty if there
 /// is no alias recorded.
 StringRef AliasState::getTypeAlias(Type ty) const {
   return typeToAlias.lookup(ty);
 }
 
 /// Print all of the referenced type aliases.
-void AliasState::printTypeAliases(raw_ostream &os) const {
+void AliasState::printTypeAliases(raw_ostream &os,
+                                  NewLineCounter &newLine) const {
   for (Type type : usedTypes) {
     auto alias = typeToAlias.find(type);
     if (alias != typeToAlias.end())
-      os << '!' << alias->second << " = type " << type << '\n';
+      os << '!' << alias->second << " = type " << type << newLine;
   }
 }
 
 /// Record a reference to the given attribute.
 void AliasState::recordAttributeReference(Attribute attr) {
   // Don't recheck attributes that have already been seen or those that
   // already have an alias.
   if (!usedAttributes.insert(attr) || attrToAlias.count(attr))

 //===----------------------------------------------------------------------===//
 // AsmState
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 namespace detail {
 class AsmStateImpl {
 public:
-  explicit AsmStateImpl(Operation *op)
-      : interfaces(op->getContext()), nameState(op, interfaces) {}
+  explicit AsmStateImpl(Operation *op, AsmState::LocationMap *locationMap)
+      : interfaces(op->getContext()), nameState(op, interfaces),
+        locationMap(locationMap) {}
 
   /// Initialize the alias state to enable the printing of aliases.
   void initializeAliases(Operation *op) {
     aliasState.initialize(op, interfaces);
   }
 
   /// Get an instance of the OpAsmDialectInterface for the given dialect, or
   /// null if one wasn't registered.
   const OpAsmDialectInterface *getOpAsmInterface(Dialect *dialect) {
     return interfaces.getInterfaceFor(dialect);
   }
 
   /// Get the state used for aliases.
   AliasState &getAliasState() { return aliasState; }
 
   /// Get the state used for SSA names.
   SSANameState &getSSANameState() { return nameState; }
 
+  /// Register the location, line and column, within the buffer that the given
+  /// operation was printed at.
+  void registerOperationLocation(Operation *op, unsigned line, unsigned col) {
+    if (locationMap)
+      (*locationMap)[op] = std::make_pair(line, col);
+  }
+
 private:
   /// Collection of OpAsm interfaces implemented in the context.
   DialectInterfaceCollection<OpAsmDialectInterface> interfaces;
 
   /// The state used for attribute and type aliases.
   AliasState aliasState;
 
   /// The state used for SSA value names.
   SSANameState nameState;
+
+  /// An optional location map to be populated.
+  AsmState::LocationMap *locationMap;
 };
 } // end namespace detail
 } // end namespace mlir
 
-AsmState::AsmState(Operation *op) : impl(std::make_unique<AsmStateImpl>(op)) {}
+AsmState::AsmState(Operation *op, LocationMap *locationMap)
+    : impl(std::make_unique<AsmStateImpl>(op, locationMap)) {}
 AsmState::~AsmState() {}
 
 //===----------------------------------------------------------------------===//
 // ModulePrinter
 //===----------------------------------------------------------------------===//
 
 namespace {
 class ModulePrinter {
protected:
   /// The output stream for the printer.
   raw_ostream &os;
 
   /// A set of flags to control the printer's behavior.
   OpPrintingFlags printerFlags;
 
   /// An optional printer state for the module.
   AsmStateImpl *state;
+
+  /// A tracker for the number of new lines emitted during printing.
+  NewLineCounter newLine;
 };
 } // end anonymous namespace
 
 void ModulePrinter::printTrailingLocation(Location loc) {
   // Check to see if we are printing debug information.
   if (!printerFlags.shouldPrintDebugInfo())
     return;
 
case StandardAttributes::CallSiteLocation: {
     if (!pretty)
       os << "callsite(";
     printLocationInternal(callee, pretty);
     if (pretty) {
       if (callee.isa<NameLoc>()) {
         if (caller.isa<FileLineColLoc>()) {
           os << " at ";
         } else {
-          os << "\n at ";
+          os << newLine << " at ";
         }
       } else {
-        os << "\n at ";
+        os << newLine << " at ";
       }
     } else {
       os << " at ";
     }
     printLocationInternal(caller, pretty);
     if (!pretty)
       os << ")";
     break;
private:
 
   // This is the current indentation level for nested structures.
   unsigned currentIndent = 0;
 };
 } // end anonymous namespace
 
 void OperationPrinter::print(ModuleOp op) {
   // Output the aliases at the top level.
-  state->getAliasState().printAttributeAliases(os);
-  state->getAliasState().printTypeAliases(os);
+  state->getAliasState().printAttributeAliases(os, newLine);
+  state->getAliasState().printTypeAliases(os, newLine);
 
   // Print the module.
   print(op.getOperation());
 }
 
 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());
 }
 
 void OperationPrinter::printOperation(Operation *op) {
   if (size_t numResults = op->getNumResults()) {
     auto printResultGroup = [&](size_t resultNo, size_t resultCount) {
if (!block->getParent()) {
       llvm::array_pod_sort(predIDs.begin(), predIDs.end());
 
       os << "\t// " << predIDs.size() << " preds: ";
 
       interleaveComma(predIDs, [&](std::pair<unsigned, Block *> pred) {
         printBlockName(pred.second);
       });
     }
-    os << '\n';
+    os << newLine;
   }
 
   currentIndent += indentWidth;
   auto range = llvm::make_range(
       block->getOperations().begin(),
       std::prev(block->getOperations().end(), printBlockTerminator ? 0 : 1));
   for (auto &op : range) {
     print(&op);
-    os << '\n';
+    os << newLine;
   }
   currentIndent -= indentWidth;
 }
 
 void OperationPrinter::printValueID(Value value, bool printResultNo) const {
   state->getSSANameState().printValueID(value, printResultNo, os);
 }
 
void OperationPrinter::printSuccessorAndUseList(Operation *term,
   os << " : ";
   interleaveComma(succOperands,
                   [this](Value operand) { printType(operand.getType()); });
   os << ')';
 }
 
 void OperationPrinter::printRegion(Region &region, bool printEntryBlockArgs,
                                    bool printBlockTerminators) {
-  os << " {\n";
+  os << " {" << newLine;
   if (!region.empty()) {
     auto *entryBlock = &region.front();
     print(entryBlock, printEntryBlockArgs && entryBlock->getNumArguments() != 0,
           printBlockTerminators);
     for (auto &b : llvm::drop_begin(region.getBlocks(), 1))
       print(&b);
   }
   os.indent(currentIndent) << "}";

mlir/lib/Transforms/CMakeLists.txt

 add_subdirectory(Utils)
 
 add_llvm_library(MLIRTransforms
   AffineDataCopyGeneration.cpp
   AffineLoopInvariantCodeMotion.cpp
   Canonicalizer.cpp
   CSE.cpp
   DialectConversion.cpp
   Inliner.cpp
+  LocationSnapshot.cpp
   LoopCoalescing.cpp
   LoopFusion.cpp
   LoopInvariantCodeMotion.cpp
   LoopTiling.cpp
   LoopUnrollAndJam.cpp
   LoopUnroll.cpp
   MemRefDataFlowOpt.cpp
   OpStats.cpp

mlir/lib/Transforms/LocationSnapshot.cpp

+//===- LocationSnapshot.cpp - Location Snapshot Utilities -----------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Transforms/LocationSnapshot.h"
+#include "mlir/IR/AsmState.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Support/FileUtilities.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/ToolOutputFile.h"
+
+using namespace mlir;
+
+/// This function generates new locations from the given IR by snapshotting the
+/// IR to the given stream, and using the printed locations within that stream.
+/// If a 'tag' is non-empty, the generated locations are represented as a
+/// NameLoc with the given tag as the name, and then fused with the existing
+/// locations. Otherwise, the existing locations are replaced.
+static void generateLocationsFromIR(raw_ostream &os, StringRef fileName,
+                                    Operation *op, OpPrintingFlags flags,
+                                    StringRef tag) {
+  // Print the IR to the stream, and collect the raw line+column information.
+  AsmState::LocationMap opToLineCol;
+  AsmState state(op, &opToLineCol);
+  op->print(os, state, flags);
+
+  Builder builder(op->getContext());
+  Optional<Identifier> tagIdentifier;
+  if (!tag.empty())
+    tagIdentifier = builder.getIdentifier(tag);
+
+  // Walk and generate new locations for each of the operations.
+  Identifier file = builder.getIdentifier(fileName);
+  op->walk([&](Operation *opIt) {
+    // Check to see if this operation has a mapped location. Some operations may
+    // be elided from the printed form, e.g. the body terminators of some region
+    // operations.
+    auto it = opToLineCol.find(opIt);
+    if (it == opToLineCol.end())
+      return;
+    const std::pair<unsigned, unsigned> &lineCol = it->second;
+    auto newLoc =
+        builder.getFileLineColLoc(file, lineCol.first, lineCol.second);
+
+    // If we don't have a tag, set the location directly
+    if (!tagIdentifier) {
+      opIt->setLoc(newLoc);
+      return;
+    }
+
+    // Otherwise, build a fused location with the existing op loc.
+    opIt->setLoc(builder.getFusedLoc(
+        {opIt->getLoc(), NameLoc::get(*tagIdentifier, newLoc)}));
+  });
+}
+
+/// This function generates new locations from the given IR by snapshotting the
+/// IR to the given file, and using the printed locations within that file. If
+/// `filename` is empty, a temporary file is generated instead.
+static LogicalResult generateLocationsFromIR(StringRef fileName, Operation *op,
+                                             OpPrintingFlags flags,
+                                             StringRef tag) {
+  // If a filename wasn't provided, then generate one.
+  SmallString<32> filepath(fileName);
+  if (filepath.empty()) {
+    if (std::error_code error = llvm::sys::fs::createTemporaryFile(
+            "mlir_snapshot", "tmp.mlir", filepath)) {
+      return op->emitError()
+             << "failed to generate temporary file for location snapshot: "
+             << error.message();
+    }
+  }
+
+  // Open the output file for emission.
+  std::string error;
+  std::unique_ptr<llvm::ToolOutputFile> outputFile =
+      openOutputFile(filepath, &error);
+  if (!outputFile)
+    return op->emitError() << error;
+
+  // Generate the intermediate locations.
+  generateLocationsFromIR(outputFile->os(), filepath, op, flags, tag);
+  outputFile->keep();
+  return success();
+}
+
+/// This function generates new locations from the given IR by snapshotting the
+/// IR to the given stream, and using the printed locations within that stream.
+/// The generated locations replace the current operation locations.
+void mlir::generateLocationsFromIR(raw_ostream &os, StringRef fileName,
+                                   Operation *op, OpPrintingFlags flags) {
+  ::generateLocationsFromIR(os, fileName, op, flags, /*tag=*/StringRef());
+}
+/// This function generates new locations from the given IR by snapshotting the
+/// IR to the given file, and using the printed locations within that file. If
+/// `filename` is empty, a temporary file is generated instead.
+LogicalResult mlir::generateLocationsFromIR(StringRef fileName, Operation *op,
+                                            OpPrintingFlags flags) {
+  return ::generateLocationsFromIR(fileName, op, flags, /*tag=*/StringRef());
+}
+
+/// This function generates new locations from the given IR by snapshotting the
+/// IR to the given stream, and using the printed locations within that stream.
+/// The generated locations are represented as a NameLoc with the given tag as
+/// the name, and then fused with the existing locations.
+void mlir::generateLocationsFromIR(raw_ostream &os, StringRef fileName,
+                                   StringRef tag, Operation *op,
+                                   OpPrintingFlags flags) {
+  ::generateLocationsFromIR(os, fileName, op, flags, tag);
+}
+/// This function generates new locations from the given IR by snapshotting the
+/// IR to the given file, and using the printed locations within that file. If
+/// `filename` is empty, a temporary file is generated instead.
+LogicalResult mlir::generateLocationsFromIR(StringRef fileName, StringRef tag,
+                                            Operation *op,
+                                            OpPrintingFlags flags) {
+  return ::generateLocationsFromIR(fileName, op, flags, tag);
+}
+
+namespace {
+class LocationSnapshotPass : public OperationPass<LocationSnapshotPass> {
+public:
+  LocationSnapshotPass() = default;
+  LocationSnapshotPass(const LocationSnapshotPass &) {}
+  LocationSnapshotPass(OpPrintingFlags flags, StringRef fileName, StringRef tag)
+      : flags(flags) {
+    this->fileName = fileName.str();
+    this->tag = tag.str();
+  }
+
+  void runOnOperation() override {
+    Operation *op = getOperation();
+    if (failed(generateLocationsFromIR(fileName, op, OpPrintingFlags(), tag)))
+      return signalPassFailure();
+  }
+
+  Option<std::string> fileName{
+      *this, "filename",
+      llvm::cl::desc("The filename to print the generated IR.")};
+  Option<std::string> tag{
+      *this, "tag",
+      llvm::cl::desc("A tag to use when fusing the new locations with the "
+                     "original. If unset, the locations are replaced.")};
+
+  /// The printing flags to use when creating the snapshot.
+  OpPrintingFlags flags;
+};
+} // end anonymous namespace
+
+std::unique_ptr<Pass> mlir::createLocationSnapshotPass(OpPrintingFlags flags,
+                                                       StringRef fileName,
+                                                       StringRef tag) {
+  return std::make_unique<LocationSnapshotPass>(flags, fileName, tag);
+}
+
+static PassRegistration<LocationSnapshotPass>
+    reg("snapshot-op-locations", "generate new locations from the current IR");

mlir/test/Transforms/location-snapshot.mlir

+// RUN: mlir-opt -snapshot-op-locations='filename=%/t' -mlir-print-debuginfo %s | FileCheck %s -DFILE=%/t
+// RUN: mlir-opt -snapshot-op-locations='filename=%/t tag='tagged'' -mlir-print-debuginfo %s | FileCheck %s --check-prefix=TAG -DFILE=%/t
+
+// CHECK-LABEL: func @function
+// CHECK-NEXT: loc("[[FILE]]":{{[0-9]+}}:{{[0-9]+}})
+// CHECK-NEXT: loc("[[FILE]]":{{[0-9]+}}:{{[0-9]+}})
+// CHECK-NEXT: } loc("[[FILE]]":{{[0-9]+}}:{{[0-9]+}})
+
+// TAG-LABEL: func @function
+// TAG-NEXT: loc(fused["original", "tagged"("[[FILE]]":{{[0-9]+}}:{{[0-9]+}})])
+// TAG-NEXT: loc(fused["original", "tagged"("[[FILE]]":{{[0-9]+}}:{{[0-9]+}})])
+// TAG-NEXT: } loc(fused["original", "tagged"("[[FILE]]":{{[0-9]+}}:{{[0-9]+}})])
+
+func @function() -> i32 {
+  %1 = "foo"() : () -> i32 loc("original")
+  return %1 : i32 loc("original")
+} loc("original")