Add support for setting the location of an IROperandOwner.

PiperOrigin-RevId: 222995814

mlir/include/mlir/IR/Instructions.h

@@ -77,6 +77,7 @@ class Instruction final
 public:
   using IROperandOwner::getContext;
   using IROperandOwner::getLoc;
+  using IROperandOwner::setLoc;
 
   //===--------------------------------------------------------------------===//
   // Operands

mlir/include/mlir/IR/Operation.h

@@ -46,6 +46,9 @@ public:
   /// The source location the operation was defined or derived from.
   Location getLoc() const;
 
+  /// Set the source location the operation was defined or derived from.
+  void setLoc(Location loc);
+
   /// Return the function this operation is defined in.  This has a verbose
   /// name to avoid name lookup ambiguities.
   Function *getOperationFunction();

mlir/include/mlir/IR/Statements.h

@@ -58,6 +58,7 @@ public:
   using Statement::getLoc;
   using Statement::moveBefore;
   using Statement::print;
+  using Statement::setLoc;
 
   /// Check if this statement is a return statement.
   bool isReturn() const;

mlir/include/mlir/IR/UseDefLists.h

@@ -90,6 +90,9 @@ public:
   /// The source location the operation was defined or derived from.
   Location getLoc() const { return locationAndKind.getPointer(); }
 
+  /// Set the source location the operation was defined or derived from.
+  void setLoc(Location loc) { locationAndKind.setPointer(loc); }
+
   /// Return the context this operation is associated with.
   MLIRContext *getContext() const;
 

mlir/lib/IR/Operation.cpp

@@ -87,6 +87,14 @@ Location Operation::getLoc() const {
   return llvm::cast<OperationStmt>(this)->getLoc();
 }
 
+/// Set the source location the operation was defined or derived from.
+void Operation::setLoc(Location loc) {
+  if (auto *inst = llvm::dyn_cast<Instruction>(this))
+    inst->setLoc(loc);
+  else
+    llvm::cast<OperationStmt>(this)->setLoc(loc);
+}
+
 /// Return the function this operation is defined in.
 Function *Operation::getOperationFunction() {
   if (auto *inst = llvm::dyn_cast<Instruction>(this))