[flang-commits] [flang] e084679 - [mlir] Make locations required when adding/creating block arguments

River Riddle via flang-commits flang-commits at lists.llvm.org
Wed Jan 19 17:36:08 PST 2022


Author: River Riddle
Date: 2022-01-19T17:35:35-08:00
New Revision: e084679f961b2a2779f7ee23d53b4d8c4df2517c

URL: https://github.com/llvm/llvm-project/commit/e084679f961b2a2779f7ee23d53b4d8c4df2517c
DIFF: https://github.com/llvm/llvm-project/commit/e084679f961b2a2779f7ee23d53b4d8c4df2517c.diff

LOG: [mlir] Make locations required when adding/creating block arguments

BlockArguments gained the ability to have locations attached a while ago, but they
have always been optional. This goes against the core tenant of MLIR where location
information is a requirement, so this commit updates the API to require locations.

Fixes #53279

Differential Revision: https://reviews.llvm.org/D117633

Added: 
    

Modified: 
    flang/lib/Optimizer/CodeGen/TargetRewrite.cpp
    flang/lib/Optimizer/Dialect/FIROps.cpp
    flang/lib/Optimizer/Transforms/AbstractResult.cpp
    flang/lib/Optimizer/Transforms/RewriteLoop.cpp
    mlir/include/mlir-c/IR.h
    mlir/include/mlir/Dialect/GPU/GPUOps.td
    mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td
    mlir/include/mlir/IR/Block.h
    mlir/include/mlir/IR/Builders.h
    mlir/include/mlir/IR/FunctionImplementation.h
    mlir/include/mlir/IR/FunctionInterfaces.h
    mlir/include/mlir/IR/FunctionInterfaces.td
    mlir/include/mlir/IR/OpImplementation.h
    mlir/include/mlir/IR/Region.h
    mlir/lib/Bindings/Python/IRCore.cpp
    mlir/lib/CAPI/Dialect/Linalg.cpp
    mlir/lib/CAPI/IR/IR.cpp
    mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
    mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp
    mlir/lib/Conversion/SCFToOpenMP/SCFToOpenMP.cpp
    mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
    mlir/lib/Conversion/SCFToStandard/SCFToStandard.cpp
    mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp
    mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
    mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
    mlir/lib/Dialect/Affine/IR/AffineOps.cpp
    mlir/lib/Dialect/Async/IR/Async.cpp
    mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp
    mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
    mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp
    mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp
    mlir/lib/Dialect/GPU/Transforms/MemoryPromotion.cpp
    mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
    mlir/lib/Dialect/LLVMIR/Transforms/LegalizeForExport.cpp
    mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp
    mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
    mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
    mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
    mlir/lib/Dialect/PDLInterp/IR/PDLInterp.cpp
    mlir/lib/Dialect/SCF/SCF.cpp
    mlir/lib/Dialect/SCF/Transforms/ForToWhile.cpp
    mlir/lib/Dialect/SCF/Transforms/Utils.cpp
    mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp
    mlir/lib/Dialect/Shape/IR/Shape.cpp
    mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
    mlir/lib/Dialect/StandardOps/IR/Ops.cpp
    mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
    mlir/lib/IR/Block.cpp
    mlir/lib/IR/Builders.cpp
    mlir/lib/IR/FunctionImplementation.cpp
    mlir/lib/IR/FunctionInterfaces.cpp
    mlir/lib/IR/Region.cpp
    mlir/lib/Parser/Parser.cpp
    mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp
    mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp
    mlir/lib/Transforms/BufferResultsToOutParams.cpp
    mlir/lib/Transforms/NormalizeMemRefs.cpp
    mlir/lib/Transforms/Utils/DialectConversion.cpp
    mlir/lib/Transforms/Utils/InliningUtils.cpp
    mlir/lib/Transforms/Utils/RegionUtils.cpp
    mlir/test/CAPI/ir.c
    mlir/test/lib/Dialect/Test/TestPatterns.cpp
    mlir/test/lib/IR/TestFunc.cpp

Removed: 
    


################################################################################
diff  --git a/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp b/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp
index 7a762fb181bf0..1e671dbb72093 100644
--- a/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp
+++ b/flang/lib/Optimizer/CodeGen/TargetRewrite.cpp
@@ -493,8 +493,8 @@ class TargetRewrite : public TargetRewriteBase<TargetRewrite> {
         case FixupTy::Codes::ArgumentAsLoad: {
           // Argument was pass-by-value, but is now pass-by-reference and
           // possibly with a 
diff erent element type.
-          auto newArg =
-              func.front().insertArgument(fixup.index, newInTys[fixup.index]);
+          auto newArg = func.front().insertArgument(fixup.index,
+                                                    newInTys[fixup.index], loc);
           rewriter->setInsertionPointToStart(&func.front());
           auto oldArgTy = ReferenceType::get(oldArgTys[fixup.index - offset]);
           auto cast = rewriter->create<ConvertOp>(loc, oldArgTy, newArg);
@@ -505,8 +505,8 @@ class TargetRewrite : public TargetRewriteBase<TargetRewrite> {
         case FixupTy::Codes::ArgumentType: {
           // Argument is pass-by-value, but its type has likely been modified to
           // suit the target ABI convention.
-          auto newArg =
-              func.front().insertArgument(fixup.index, newInTys[fixup.index]);
+          auto newArg = func.front().insertArgument(fixup.index,
+                                                    newInTys[fixup.index], loc);
           rewriter->setInsertionPointToStart(&func.front());
           auto mem =
               rewriter->create<fir::AllocaOp>(loc, newInTys[fixup.index]);
@@ -524,8 +524,8 @@ class TargetRewrite : public TargetRewriteBase<TargetRewrite> {
         case FixupTy::Codes::CharPair: {
           // The FIR boxchar argument has been split into a pair of distinct
           // arguments that are in juxtaposition to each other.
-          auto newArg =
-              func.front().insertArgument(fixup.index, newInTys[fixup.index]);
+          auto newArg = func.front().insertArgument(fixup.index,
+                                                    newInTys[fixup.index], loc);
           if (fixup.second == 1) {
             rewriter->setInsertionPointToStart(&func.front());
             auto boxTy = oldArgTys[fixup.index - offset - fixup.second];
@@ -539,8 +539,8 @@ class TargetRewrite : public TargetRewriteBase<TargetRewrite> {
         case FixupTy::Codes::ReturnAsStore: {
           // The value being returned is now being returned in memory (callee
           // stack space) through a hidden reference argument.
-          auto newArg =
-              func.front().insertArgument(fixup.index, newInTys[fixup.index]);
+          auto newArg = func.front().insertArgument(fixup.index,
+                                                    newInTys[fixup.index], loc);
           offset++;
           func.walk([&](mlir::ReturnOp ret) {
             rewriter->setInsertionPoint(ret);
@@ -571,8 +571,8 @@ class TargetRewrite : public TargetRewriteBase<TargetRewrite> {
         case FixupTy::Codes::Split: {
           // The FIR argument has been split into a pair of distinct arguments
           // that are in juxtaposition to each other. (For COMPLEX value.)
-          auto newArg =
-              func.front().insertArgument(fixup.index, newInTys[fixup.index]);
+          auto newArg = func.front().insertArgument(fixup.index,
+                                                    newInTys[fixup.index], loc);
           if (fixup.second == 1) {
             rewriter->setInsertionPointToStart(&func.front());
             auto cplxTy = oldArgTys[fixup.index - offset - fixup.second];
@@ -594,9 +594,10 @@ class TargetRewrite : public TargetRewriteBase<TargetRewrite> {
           // The first part of the pair appears in the original argument
           // position. The second part of the pair is appended after all the
           // original arguments. (Boxchar arguments.)
-          auto newBufArg =
-              func.front().insertArgument(fixup.index, newInTys[fixup.index]);
-          auto newLenArg = func.front().addArgument(trailingTys[fixup.second]);
+          auto newBufArg = func.front().insertArgument(
+              fixup.index, newInTys[fixup.index], loc);
+          auto newLenArg =
+              func.front().addArgument(trailingTys[fixup.second], loc);
           auto boxTy = oldArgTys[fixup.index - offset];
           rewriter->setInsertionPointToStart(&func.front());
           auto box =

diff  --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp
index 852abefcd1d11..36c7c8f57461d 100644
--- a/flang/lib/Optimizer/Dialect/FIROps.cpp
+++ b/flang/lib/Optimizer/Dialect/FIROps.cpp
@@ -1534,9 +1534,11 @@ void fir::IterWhileOp::build(mlir::OpBuilder &builder,
     result.addTypes(v.getType());
   mlir::Region *bodyRegion = result.addRegion();
   bodyRegion->push_back(new Block{});
-  bodyRegion->front().addArgument(builder.getIndexType());
-  bodyRegion->front().addArgument(iterate.getType());
-  bodyRegion->front().addArguments(iterArgs.getTypes());
+  bodyRegion->front().addArgument(builder.getIndexType(), result.location);
+  bodyRegion->front().addArgument(iterate.getType(), result.location);
+  bodyRegion->front().addArguments(
+      iterArgs.getTypes(),
+      SmallVector<Location>(iterArgs.size(), result.location));
   result.addAttributes(attributes);
 }
 
@@ -1854,8 +1856,10 @@ void fir::DoLoopOp::build(mlir::OpBuilder &builder,
   bodyRegion->push_back(new Block{});
   if (iterArgs.empty() && !finalCountValue)
     DoLoopOp::ensureTerminator(*bodyRegion, builder, result.location);
-  bodyRegion->front().addArgument(builder.getIndexType());
-  bodyRegion->front().addArguments(iterArgs.getTypes());
+  bodyRegion->front().addArgument(builder.getIndexType(), result.location);
+  bodyRegion->front().addArguments(
+      iterArgs.getTypes(),
+      SmallVector<Location>(iterArgs.size(), result.location));
   if (unordered)
     result.addAttribute(unorderedAttrName(result.name), builder.getUnitAttr());
   result.addAttributes(attributes);

diff  --git a/flang/lib/Optimizer/Transforms/AbstractResult.cpp b/flang/lib/Optimizer/Transforms/AbstractResult.cpp
index 3ce355e6ec25e..465755c87ccab 100644
--- a/flang/lib/Optimizer/Transforms/AbstractResult.cpp
+++ b/flang/lib/Optimizer/Transforms/AbstractResult.cpp
@@ -231,7 +231,7 @@ class AbstractResultOpt : public fir::AbstractResultOptBase<AbstractResultOpt> {
         mlir::OpBuilder rewriter(context);
         auto resultType = funcTy.getResult(0);
         auto argTy = getResultArgumentType(resultType, options);
-        options.newArg = func.front().insertArgument(zero, argTy);
+        options.newArg = func.front().insertArgument(zero, argTy, loc);
         if (mustEmboxResult(resultType, options)) {
           auto bufferType = fir::ReferenceType::get(resultType);
           rewriter.setInsertionPointToStart(&func.front());

diff  --git a/flang/lib/Optimizer/Transforms/RewriteLoop.cpp b/flang/lib/Optimizer/Transforms/RewriteLoop.cpp
index 5f49be5f46e9e..bd13a1dc98e0c 100644
--- a/flang/lib/Optimizer/Transforms/RewriteLoop.cpp
+++ b/flang/lib/Optimizer/Transforms/RewriteLoop.cpp
@@ -49,7 +49,7 @@ class CfgLoopConv : public mlir::OpRewritePattern<fir::DoLoopOp> {
     // conditional block since it already has the induction variable and
     // loop-carried values as arguments.
     auto *conditionalBlock = &loop.region().front();
-    conditionalBlock->addArgument(rewriter.getIndexType());
+    conditionalBlock->addArgument(rewriter.getIndexType(), loc);
     auto *firstBlock =
         rewriter.splitBlock(conditionalBlock, conditionalBlock->begin());
     auto *lastBlock = &loop.region().back();

diff  --git a/mlir/include/mlir-c/IR.h b/mlir/include/mlir-c/IR.h
index 1d884e634b2f2..e5c20ae70cacf 100644
--- a/mlir/include/mlir-c/IR.h
+++ b/mlir/include/mlir-c/IR.h
@@ -531,7 +531,8 @@ MLIR_CAPI_EXPORTED MlirRegion mlirRegionGetNextInOperation(MlirRegion region);
 /// Creates a new empty block with the given argument types and transfers
 /// ownership to the caller.
 MLIR_CAPI_EXPORTED MlirBlock mlirBlockCreate(intptr_t nArgs,
-                                             MlirType const *args);
+                                             MlirType const *args,
+                                             MlirLocation const *locs);
 
 /// Takes a block owned by the caller and destroys it.
 MLIR_CAPI_EXPORTED void mlirBlockDestroy(MlirBlock block);
@@ -590,7 +591,8 @@ MLIR_CAPI_EXPORTED intptr_t mlirBlockGetNumArguments(MlirBlock block);
 /// Appends an argument of the specified type to the block. Returns the newly
 /// added argument.
 MLIR_CAPI_EXPORTED MlirValue mlirBlockAddArgument(MlirBlock block,
-                                                  MlirType type);
+                                                  MlirType type,
+                                                  MlirLocation loc);
 
 /// Returns `pos`-th argument of the block.
 MLIR_CAPI_EXPORTED MlirValue mlirBlockGetArgument(MlirBlock block,

diff  --git a/mlir/include/mlir/Dialect/GPU/GPUOps.td b/mlir/include/mlir/Dialect/GPU/GPUOps.td
index f48914df9ef5a..4e54bda5ab0eb 100644
--- a/mlir/include/mlir/Dialect/GPU/GPUOps.td
+++ b/mlir/include/mlir/Dialect/GPU/GPUOps.td
@@ -246,7 +246,7 @@ def GPU_GPUFuncOp : GPU_Op<"func", [
 
     /// Adds a new block argument that corresponds to buffers located in
     /// workgroup memory.
-    BlockArgument addWorkgroupAttribution(Type type);
+    BlockArgument addWorkgroupAttribution(Type type, Location loc);
 
     /// Returns the number of buffers located in the private memory.
     unsigned getNumPrivateAttributions() {
@@ -267,7 +267,7 @@ def GPU_GPUFuncOp : GPU_Op<"func", [
 
     /// Adds a new block argument that corresponds to buffers located in
     /// private memory.
-    BlockArgument addPrivateAttribution(Type type);
+    BlockArgument addPrivateAttribution(Type type, Location loc);
 
     /// Returns the name of the attribute containing the number of buffers
     /// located in the workgroup memory.

diff  --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td
index 631f7ca16fbde..4b8ae4985ca53 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td
@@ -545,7 +545,8 @@ def Linalg_TiledLoopOp : Linalg_Op<"tiled_loop", [
 
       getOperation()->insertOperands(getNumControlOperands() + numInputs,
                                      operand);
-      getBody()->insertArgument(numLoops + numInputs, operand.getType());
+      getBody()->insertArgument(numLoops + numInputs, operand.getType(), 
+                                getLoc());
       getOperation()->setAttr(
           TiledLoopOp::getOperandSegmentSizeAttr(),
           builder.getI32VectorAttr(
@@ -562,7 +563,7 @@ def Linalg_TiledLoopOp : Linalg_Op<"tiled_loop", [
       getOperation()->insertOperands(
           getNumControlOperands() + numInputs + numOutputs, operand);
       getBody()->insertArgument(numLoops + numInputs + numOutputs,
-                                operand.getType());
+                                operand.getType(), getLoc());
       getOperation()->setAttr(
           TiledLoopOp::getOperandSegmentSizeAttr(),
           builder.getI32VectorAttr(

diff  --git a/mlir/include/mlir/IR/Block.h b/mlir/include/mlir/IR/Block.h
index 623020764fc3b..3f5faa74d3302 100644
--- a/mlir/include/mlir/IR/Block.h
+++ b/mlir/include/mlir/IR/Block.h
@@ -88,21 +88,20 @@ class Block : public IRObjectWithUseList<BlockOperand>,
   bool args_empty() { return arguments.empty(); }
 
   /// Add one value to the argument list.
-  BlockArgument addArgument(Type type, Optional<Location> loc = {});
+  BlockArgument addArgument(Type type, 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,
-                               Optional<Location> loc = {});
+  BlockArgument insertArgument(args_iterator it, Type type, Location loc);
 
   /// Add one argument to the argument list for each type specified in the list.
+  /// `locs` is required to have the same number of elements as `types`.
   iterator_range<args_iterator> addArguments(TypeRange types,
-                                             ArrayRef<Location> locs = {});
+                                             ArrayRef<Location> locs);
 
   /// Add one value to the argument list at the specified position.
-  BlockArgument insertArgument(unsigned index, Type type,
-                               Optional<Location> loc = {});
+  BlockArgument insertArgument(unsigned index, Type type, Location loc);
 
   /// Erase the argument at 'index' and remove it from the argument list.
   void eraseArgument(unsigned index);

diff  --git a/mlir/include/mlir/IR/Builders.h b/mlir/include/mlir/IR/Builders.h
index cf4cdb9203e53..a78ec406bc9d4 100644
--- a/mlir/include/mlir/IR/Builders.h
+++ b/mlir/include/mlir/IR/Builders.h
@@ -384,15 +384,18 @@ class OpBuilder : public Builder {
 
   /// 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'.
+  /// 'parent'. `locs` contains the locations of the inserted arguments, and
+  /// should match the size of `argTypes`.
   Block *createBlock(Region *parent, Region::iterator insertPt = {},
                      TypeRange argTypes = llvm::None,
-                     ArrayRef<Location> locs = {});
+                     ArrayRef<Location> locs = llvm::None);
 
   /// Add new block with 'argTypes' arguments and set the insertion point to the
-  /// end of it. The block is placed before 'insertBefore'.
+  /// end of it. The block is placed before 'insertBefore'. `locs` contains the
+  /// locations of the inserted arguments, and should match the size of
+  /// `argTypes`.
   Block *createBlock(Block *insertBefore, TypeRange argTypes = llvm::None,
-                     ArrayRef<Location> locs = {});
+                     ArrayRef<Location> locs = llvm::None);
 
   //===--------------------------------------------------------------------===//
   // Operation Creation

diff  --git a/mlir/include/mlir/IR/FunctionImplementation.h b/mlir/include/mlir/IR/FunctionImplementation.h
index d1a15a78e5645..3562c6ec1537c 100644
--- 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::OperandType> &argNames,
     SmallVectorImpl<Type> &argTypes, SmallVectorImpl<NamedAttrList> &argAttrs,
-    SmallVectorImpl<Optional<Location>> &argLocations, bool &isVariadic);
+    SmallVectorImpl<Location> &argLocations, bool &isVariadic);
 
 /// Parses a function signature using `parser`. The `allowVariadic` argument
 /// indicates whether functions with variadic arguments are supported. The
@@ -71,7 +71,7 @@ parseFunctionSignature(OpAsmParser &parser, bool allowVariadic,
                        SmallVectorImpl<OpAsmParser::OperandType> &argNames,
                        SmallVectorImpl<Type> &argTypes,
                        SmallVectorImpl<NamedAttrList> &argAttrs,
-                       SmallVectorImpl<Optional<Location>> &argLocations,
+                       SmallVectorImpl<Location> &argLocations,
                        bool &isVariadic, SmallVectorImpl<Type> &resultTypes,
                        SmallVectorImpl<NamedAttrList> &resultAttrs);
 

diff  --git a/mlir/include/mlir/IR/FunctionInterfaces.h b/mlir/include/mlir/IR/FunctionInterfaces.h
index 0c99776fdee1e..b81abceef6149 100644
--- a/mlir/include/mlir/IR/FunctionInterfaces.h
+++ b/mlir/include/mlir/IR/FunctionInterfaces.h
@@ -72,7 +72,7 @@ inline ArrayRef<NamedAttribute> getResultAttrs(Operation *op, unsigned index) {
 void insertFunctionArguments(Operation *op, ArrayRef<unsigned> argIndices,
                              TypeRange argTypes,
                              ArrayRef<DictionaryAttr> argAttrs,
-                             ArrayRef<Optional<Location>> argLocs,
+                             ArrayRef<Location> argLocs,
                              unsigned originalNumArgs, Type newType);
 
 /// Insert the specified results and update the function type attribute.

diff  --git a/mlir/include/mlir/IR/FunctionInterfaces.td b/mlir/include/mlir/IR/FunctionInterfaces.td
index 17f5a5f99f957..b6993ff6ffc76 100644
--- a/mlir/include/mlir/IR/FunctionInterfaces.td
+++ b/mlir/include/mlir/IR/FunctionInterfaces.td
@@ -156,7 +156,13 @@ def FunctionOpInterface : OpInterface<"FunctionOpInterface"> {
       assert(empty() && "function already has an entry block");
       Block *entry = new Block();
       push_back(entry);
-      entry->addArguments($_op.getArgumentTypes());
+      
+      // FIXME: Allow for passing in locations for these arguments instead of using
+      // the operations location.
+      ArrayRef<Type> inputTypes = $_op.getArgumentTypes();
+      SmallVector<Location> locations(inputTypes.size(),
+                                      $_op.getOperation()->getLoc());
+      entry->addArguments(inputTypes, locations);
       return entry;
     }
 
@@ -234,7 +240,7 @@ def FunctionOpInterface : OpInterface<"FunctionOpInterface"> {
     /// Insert a single argument of type `argType` with attributes `argAttrs` and
     /// location `argLoc` at `argIndex`.
     void insertArgument(unsigned argIndex, Type argType, DictionaryAttr argAttrs,
-                        Optional<Location> argLoc = {}) {
+                        Location argLoc) {
       insertArguments({argIndex}, {argType}, {argAttrs}, {argLoc});
     }
 
@@ -244,7 +250,7 @@ def FunctionOpInterface : OpInterface<"FunctionOpInterface"> {
     /// appear in the same order that they were listed here.
     void insertArguments(ArrayRef<unsigned> argIndices, TypeRange argTypes,
                         ArrayRef<DictionaryAttr> argAttrs,
-                        ArrayRef<Optional<Location>> argLocs) {
+                        ArrayRef<Location> argLocs) {
       unsigned originalNumArgs = $_op.getNumArguments();
       Type newType = $_op.getTypeWithArgsAndResults(
           argIndices, argTypes, /*resultIndices=*/{}, /*resultTypes=*/{});

diff  --git a/mlir/include/mlir/IR/OpImplementation.h b/mlir/include/mlir/IR/OpImplementation.h
index 2f8cea6303275..56e7c72093219 100644
--- a/mlir/include/mlir/IR/OpImplementation.h
+++ b/mlir/include/mlir/IR/OpImplementation.h
@@ -1209,22 +1209,22 @@ class OpAsmParser : public AsmParser {
   /// 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 an optional location to be attached to each
-  /// argument. If 'enableNameShadowing' is set to true, the argument names are
-  /// allowed to shadow the names of other existing SSA values defined above the
-  /// region scope. 'enableNameShadowing' can only be set to true for regions
-  /// attached to operations that are 'IsolatedFromAbove'.
-  virtual ParseResult
-  parseRegion(Region &region, ArrayRef<OperandType> arguments = {},
-              ArrayRef<Type> argTypes = {},
-              ArrayRef<Optional<Location>> argLocations = {},
-              bool enableNameShadowing = false) = 0;
+  /// 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'.
+  virtual ParseResult parseRegion(Region &region,
+                                  ArrayRef<OperandType> arguments = {},
+                                  ArrayRef<Type> argTypes = {},
+                                  ArrayRef<Location> argLocations = {},
+                                  bool enableNameShadowing = false) = 0;
 
   /// Parses a region if present.
   virtual OptionalParseResult
   parseOptionalRegion(Region &region, ArrayRef<OperandType> arguments = {},
                       ArrayRef<Type> argTypes = {},
-                      ArrayRef<Optional<Location>> argLocations = {},
+                      ArrayRef<Location> argLocations = {},
                       bool enableNameShadowing = false) = 0;
 
   /// Parses a region if present. If the region is present, a new region is

diff  --git a/mlir/include/mlir/IR/Region.h b/mlir/include/mlir/IR/Region.h
index 44ede7b55b6b7..c2694ad33ddac 100644
--- a/mlir/include/mlir/IR/Region.h
+++ b/mlir/include/mlir/IR/Region.h
@@ -82,6 +82,7 @@ class Region {
     return empty() ? BlockArgListType() : front().getArguments();
   }
 
+  /// Returns the argument types of the first block within the region.
   ValueTypeRange<BlockArgListType> getArgumentTypes();
 
   using args_iterator = BlockArgListType::iterator;
@@ -94,21 +95,26 @@ class Region {
   bool args_empty() { return getArguments().empty(); }
 
   /// Add one value to the argument list.
-  BlockArgument addArgument(Type type) { return front().addArgument(type); }
+  BlockArgument addArgument(Type type, Location loc) {
+    return front().addArgument(type, 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) {
-    return front().insertArgument(it, type);
+  BlockArgument insertArgument(args_iterator it, Type type, Location loc) {
+    return front().insertArgument(it, type, loc);
   }
 
   /// Add one argument to the argument list for each type specified in the list.
-  iterator_range<args_iterator> addArguments(TypeRange types);
+  /// `locs` contains the locations for each of the new arguments, and must be
+  /// of equal size to `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) {
-    return front().insertArgument(index, type);
+  BlockArgument insertArgument(unsigned index, Type type, Location loc) {
+    return front().insertArgument(index, type, loc);
   }
 
   /// Erase the argument at 'index' and remove it from the argument list.

diff  --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 729431b98ec32..621c095021c7f 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -325,12 +325,18 @@ class PyBlockList {
   PyBlock appendBlock(const py::args &pyArgTypes) {
     operation->checkValid();
     llvm::SmallVector<MlirType, 4> argTypes;
+    llvm::SmallVector<MlirLocation, 4> argLocs;
     argTypes.reserve(pyArgTypes.size());
+    argLocs.reserve(pyArgTypes.size());
     for (auto &pyArg : pyArgTypes) {
       argTypes.push_back(pyArg.cast<PyType &>());
+      // TODO: Pass in a proper location here.
+      argLocs.push_back(
+          mlirLocationUnknownGet(mlirTypeGetContext(argTypes.back())));
     }
 
-    MlirBlock block = mlirBlockCreate(argTypes.size(), argTypes.data());
+    MlirBlock block =
+        mlirBlockCreate(argTypes.size(), argTypes.data(), argLocs.data());
     mlirRegionAppendOwnedBlock(region, block);
     return PyBlock(operation, block);
   }
@@ -2717,12 +2723,18 @@ void mlir::python::populateIRCore(py::module &m) {
           [](PyRegion &parent, py::list pyArgTypes) {
             parent.checkValid();
             llvm::SmallVector<MlirType, 4> argTypes;
+            llvm::SmallVector<MlirLocation, 4> argLocs;
             argTypes.reserve(pyArgTypes.size());
+            argLocs.reserve(pyArgTypes.size());
             for (auto &pyArg : pyArgTypes) {
               argTypes.push_back(pyArg.cast<PyType &>());
+              // TODO: Pass in a proper location here.
+              argLocs.push_back(
+                  mlirLocationUnknownGet(mlirTypeGetContext(argTypes.back())));
             }
 
-            MlirBlock block = mlirBlockCreate(argTypes.size(), argTypes.data());
+            MlirBlock block = mlirBlockCreate(argTypes.size(), argTypes.data(),
+                                              argLocs.data());
             mlirRegionInsertOwnedBlock(parent, 0, block);
             return PyBlock(parent.getParentOperation(), block);
           },
@@ -2734,12 +2746,18 @@ void mlir::python::populateIRCore(py::module &m) {
           [](PyBlock &self, py::args pyArgTypes) {
             self.checkValid();
             llvm::SmallVector<MlirType, 4> argTypes;
+            llvm::SmallVector<MlirLocation, 4> argLocs;
             argTypes.reserve(pyArgTypes.size());
+            argLocs.reserve(pyArgTypes.size());
             for (auto &pyArg : pyArgTypes) {
               argTypes.push_back(pyArg.cast<PyType &>());
+              // TODO: Pass in a proper location here.
+              argLocs.push_back(
+                  mlirLocationUnknownGet(mlirTypeGetContext(argTypes.back())));
             }
 
-            MlirBlock block = mlirBlockCreate(argTypes.size(), argTypes.data());
+            MlirBlock block = mlirBlockCreate(argTypes.size(), argTypes.data(),
+                                              argLocs.data());
             MlirRegion region = mlirBlockGetParentRegion(self.get());
             mlirRegionInsertOwnedBlockBefore(region, self.get(), block);
             return PyBlock(self.getParentOperation(), block);
@@ -2751,12 +2769,18 @@ void mlir::python::populateIRCore(py::module &m) {
           [](PyBlock &self, py::args pyArgTypes) {
             self.checkValid();
             llvm::SmallVector<MlirType, 4> argTypes;
+            llvm::SmallVector<MlirLocation, 4> argLocs;
             argTypes.reserve(pyArgTypes.size());
+            argLocs.reserve(pyArgTypes.size());
             for (auto &pyArg : pyArgTypes) {
               argTypes.push_back(pyArg.cast<PyType &>());
-            }
 
-            MlirBlock block = mlirBlockCreate(argTypes.size(), argTypes.data());
+              // TODO: Pass in a proper location here.
+              argLocs.push_back(
+                  mlirLocationUnknownGet(mlirTypeGetContext(argTypes.back())));
+            }
+            MlirBlock block = mlirBlockCreate(argTypes.size(), argTypes.data(),
+                                              argLocs.data());
             MlirRegion region = mlirBlockGetParentRegion(self.get());
             mlirRegionInsertOwnedBlockAfter(region, self.get(), block);
             return PyBlock(self.getParentOperation(), block);

diff  --git a/mlir/lib/CAPI/Dialect/Linalg.cpp b/mlir/lib/CAPI/Dialect/Linalg.cpp
index 6c5ba9a883811..8862b6b154ea5 100644
--- a/mlir/lib/CAPI/Dialect/Linalg.cpp
+++ b/mlir/lib/CAPI/Dialect/Linalg.cpp
@@ -28,12 +28,15 @@ void mlirLinalgFillBuiltinNamedOpRegion(MlirOperation mlirOp) {
          "Expected Linalg op with 0 blocks");
 
   SmallVector<Type, 8> argTypes;
-  for (OpOperand *opOperand : linalgOp.getInputAndOutputOperands())
+  SmallVector<Location, 8> argLocs;
+  for (OpOperand *opOperand : linalgOp.getInputAndOutputOperands()) {
     argTypes.push_back(getElementTypeOrSelf(opOperand->get().getType()));
+    argLocs.push_back(opOperand->get().getLoc());
+  }
 
   ImplicitLocOpBuilder b(op->getLoc(), op->getContext());
   Region &region = op->getRegion(0);
-  Block *body = b.createBlock(&region, /*insertPt=*/{}, argTypes);
+  Block *body = b.createBlock(&region, /*insertPt=*/{}, argTypes, argLocs);
   b.setInsertionPointToStart(body);
   fun(b, *body);
 }

diff  --git a/mlir/lib/CAPI/IR/IR.cpp b/mlir/lib/CAPI/IR/IR.cpp
index 13b7673e2bccd..9b60b11fe463e 100644
--- a/mlir/lib/CAPI/IR/IR.cpp
+++ b/mlir/lib/CAPI/IR/IR.cpp
@@ -535,10 +535,11 @@ void mlirRegionDestroy(MlirRegion region) {
 // Block API.
 //===----------------------------------------------------------------------===//
 
-MlirBlock mlirBlockCreate(intptr_t nArgs, MlirType const *args) {
+MlirBlock mlirBlockCreate(intptr_t nArgs, MlirType const *args,
+                          MlirLocation const *locs) {
   Block *b = new Block;
   for (intptr_t i = 0; i < nArgs; ++i)
-    b->addArgument(unwrap(args[i]));
+    b->addArgument(unwrap(args[i]), unwrap(locs[i]));
   return wrap(b);
 }
 
@@ -618,8 +619,9 @@ intptr_t mlirBlockGetNumArguments(MlirBlock block) {
   return static_cast<intptr_t>(unwrap(block)->getNumArguments());
 }
 
-MlirValue mlirBlockAddArgument(MlirBlock block, MlirType type) {
-  return wrap(unwrap(block)->addArgument(unwrap(type)));
+MlirValue mlirBlockAddArgument(MlirBlock block, MlirType type,
+                               MlirLocation loc) {
+  return wrap(unwrap(block)->addArgument(unwrap(type), unwrap(loc)));
 }
 
 MlirValue mlirBlockGetArgument(MlirBlock block, intptr_t pos) {

diff  --git a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
index 72eeca1662221..1a3384ca2c040 100644
--- a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
+++ b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
@@ -208,8 +208,10 @@ struct AllocaScopeOpLowering
     if (allocaScopeOp.getNumResults() == 0) {
       continueBlock = remainingOpsBlock;
     } else {
-      continueBlock = rewriter.createBlock(remainingOpsBlock,
-                                           allocaScopeOp.getResultTypes());
+      continueBlock = rewriter.createBlock(
+          remainingOpsBlock, allocaScopeOp.getResultTypes(),
+          SmallVector<Location>(allocaScopeOp->getNumResults(),
+                                allocaScopeOp.getLoc()));
       rewriter.create<LLVM::BrOp>(loc, ValueRange(), remainingOpsBlock);
     }
 
@@ -1019,7 +1021,7 @@ struct MemRefReshapeOpLowering
     Block::iterator remainingOpsIt = std::next(rewriter.getInsertionPoint());
 
     Block *condBlock = rewriter.createBlock(initBlock->getParent(), {},
-                                            {indexType, indexType});
+                                            {indexType, indexType}, {loc, loc});
 
     // Move the remaining initBlock ops to condBlock.
     Block *remainingBlock = rewriter.splitBlock(initBlock, remainingOpsIt);

diff  --git a/mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp b/mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp
index ff8e1bbaf04f3..5f6c6fc8d0815 100644
--- a/mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp
+++ b/mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp
@@ -628,7 +628,8 @@ SymbolRefAttr PatternLowering::generateRewriter(
     Position *inputPos = valueToPosition.lookup(oldValue);
     assert(inputPos && "expected value to be a pattern input");
     usedMatchValues.push_back(inputPos);
-    return newValue = rewriterFunc.front().addArgument(oldValue.getType());
+    return newValue = rewriterFunc.front().addArgument(oldValue.getType(),
+                                                       oldValue.getLoc());
   };
 
   // If this is a custom rewriter, simply dispatch to the registered rewrite

diff  --git a/mlir/lib/Conversion/SCFToOpenMP/SCFToOpenMP.cpp b/mlir/lib/Conversion/SCFToOpenMP/SCFToOpenMP.cpp
index 14ac92e07c47a..e472af257776b 100644
--- a/mlir/lib/Conversion/SCFToOpenMP/SCFToOpenMP.cpp
+++ b/mlir/lib/Conversion/SCFToOpenMP/SCFToOpenMP.cpp
@@ -188,7 +188,7 @@ static omp::ReductionDeclareOp createDecl(PatternRewriter &builder,
 
   Type type = reduce.getOperand().getType();
   builder.createBlock(&decl.initializerRegion(), decl.initializerRegion().end(),
-                      {type});
+                      {type}, {reduce.getOperand().getLoc()});
   builder.setInsertionPointToEnd(&decl.initializerRegion().back());
   Value init =
       builder.create<LLVM::ConstantOp>(reduce.getLoc(), type, initValue);
@@ -214,8 +214,10 @@ static omp::ReductionDeclareOp addAtomicRMW(OpBuilder &builder,
   OpBuilder::InsertionGuard guard(builder);
   Type type = reduce.getOperand().getType();
   Type ptrType = LLVM::LLVMPointerType::get(type);
+  Location reduceOperandLoc = reduce.getOperand().getLoc();
   builder.createBlock(&decl.atomicReductionRegion(),
-                      decl.atomicReductionRegion().end(), {ptrType, ptrType});
+                      decl.atomicReductionRegion().end(), {ptrType, ptrType},
+                      {reduceOperandLoc, reduceOperandLoc});
   Block *atomicBlock = &decl.atomicReductionRegion().back();
   builder.setInsertionPointToEnd(atomicBlock);
   Value loaded = builder.create<LLVM::LoadOp>(reduce.getLoc(),

diff  --git a/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp b/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
index 6bb3da666ce77..3a67428cc527d 100644
--- a/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
+++ b/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
@@ -177,10 +177,11 @@ ForOpConversion::matchAndRewrite(scf::ForOp forOp, OpAdaptor adaptor,
   loopOp.body().getBlocks().insert(getBlockIt(loopOp.body(), 1), header);
 
   // Create the new induction variable to use.
+  Value adapLowerBound = adaptor.getLowerBound();
   BlockArgument newIndVar =
-      header->addArgument(adaptor.getLowerBound().getType());
+      header->addArgument(adapLowerBound.getType(), adapLowerBound.getLoc());
   for (Value arg : adaptor.getInitArgs())
-    header->addArgument(arg.getType());
+    header->addArgument(arg.getType(), arg.getLoc());
   Block *body = forOp.getBody();
 
   // Apply signature conversion to the body of the forOp. It has a single block,

diff  --git a/mlir/lib/Conversion/SCFToStandard/SCFToStandard.cpp b/mlir/lib/Conversion/SCFToStandard/SCFToStandard.cpp
index 8ceaa864721c0..31e4cca72e90b 100644
--- a/mlir/lib/Conversion/SCFToStandard/SCFToStandard.cpp
+++ b/mlir/lib/Conversion/SCFToStandard/SCFToStandard.cpp
@@ -368,7 +368,8 @@ LogicalResult IfLowering::matchAndRewrite(IfOp ifOp,
     continueBlock = remainingOpsBlock;
   } else {
     continueBlock =
-        rewriter.createBlock(remainingOpsBlock, ifOp.getResultTypes());
+        rewriter.createBlock(remainingOpsBlock, ifOp.getResultTypes(),
+                             SmallVector<Location>(ifOp.getNumResults(), loc));
     rewriter.create<BranchOp>(loc, remainingOpsBlock);
   }
 
@@ -433,9 +434,10 @@ ExecuteRegionLowering::matchAndRewrite(ExecuteRegionOp op,
   rewriter.inlineRegionBefore(region, remainingOpsBlock);
 
   SmallVector<Value> vals;
-  for (auto arg : remainingOpsBlock->addArguments(op->getResultTypes())) {
+  SmallVector<Location> argLocs(op.getNumResults(), op->getLoc());
+  for (auto arg :
+       remainingOpsBlock->addArguments(op->getResultTypes(), argLocs))
     vals.push_back(arg);
-  }
   rewriter.replaceOp(op, vals);
   return success();
 }

diff  --git a/mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp b/mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp
index c43ed30a76fac..f914bad4c6eaf 100644
--- a/mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp
+++ b/mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp
@@ -812,9 +812,9 @@ struct GenericAtomicRMWOpLowering
 
     // Split the block into initial, loop, and ending parts.
     auto *initBlock = rewriter.getInsertionBlock();
-    auto *loopBlock =
-        rewriter.createBlock(initBlock->getParent(),
-                             std::next(Region::iterator(initBlock)), valueType);
+    auto *loopBlock = rewriter.createBlock(
+        initBlock->getParent(), std::next(Region::iterator(initBlock)),
+        valueType, loc);
     auto *endBlock = rewriter.createBlock(
         loopBlock->getParent(), std::next(Region::iterator(loopBlock)));
 

diff  --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
index 5e94633a7b60a..0f53b30125625 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
@@ -267,10 +267,13 @@ createLinalgBodyCalculationForElementwiseOp(Operation *op, ValueRange args,
 
     SmallVector<Value> operands = {args[0], leadingZeros, zero};
     SmallVector<Type> types = {elementTy, elementTy, elementTy};
+    SmallVector<Location> locations = {loc, loc, loc};
 
     auto whileOp = rewriter.create<scf::WhileOp>(loc, types, operands);
-    Block *before = rewriter.createBlock(&whileOp.getBefore(), {}, types);
-    Block *after = rewriter.createBlock(&whileOp.getAfter(), {}, types);
+    Block *before =
+        rewriter.createBlock(&whileOp.getBefore(), {}, types, locations);
+    Block *after =
+        rewriter.createBlock(&whileOp.getAfter(), {}, types, locations);
 
     // The conditional block of the while loop.
     {
@@ -1409,7 +1412,7 @@ class ResizeConverter : public OpRewritePattern<tosa::ResizeOp> {
 
     OpBuilder::InsertionGuard regionGuard(rewriter);
     rewriter.createBlock(&genericOp.region(), genericOp.region().end(),
-                         TypeRange({resultElementTy}));
+                         TypeRange({resultElementTy}), loc);
     Value batch = rewriter.create<linalg::IndexOp>(loc, 0);
     Value y = rewriter.create<linalg::IndexOp>(loc, 1);
     Value x = rewriter.create<linalg::IndexOp>(loc, 2);
@@ -2159,9 +2162,9 @@ class TableConverter : public OpRewritePattern<tosa::TableOp> {
 
     {
       OpBuilder::InsertionGuard regionGuard(rewriter);
-      Block *block =
-          rewriter.createBlock(&genericOp.region(), genericOp.region().end(),
-                               TypeRange({inputElementTy, resultElementTy}));
+      Block *block = rewriter.createBlock(
+          &genericOp.region(), genericOp.region().end(),
+          TypeRange({inputElementTy, resultElementTy}), {loc, loc});
 
       auto inputValue = block->getArgument(0);
       rewriter.setInsertionPointToStart(block);

diff  --git a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
index dcba23cb2ef18..6b513c4635e4a 100644
--- a/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
+++ b/mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
@@ -435,8 +435,8 @@ static scf::ForOp replaceForOpWithNewSignature(OpBuilder &b, scf::ForOp loop,
   newLoop.getLoopBody().getBlocks().splice(
       newLoop.getLoopBody().getBlocks().begin(),
       loop.getLoopBody().getBlocks());
-  for (auto operand : newIterOperands)
-    newLoop.getBody()->addArgument(operand.getType());
+  for (Value operand : newIterOperands)
+    newLoop.getBody()->addArgument(operand.getType(), operand.getLoc());
 
   for (auto it : llvm::zip(loop.getResults(), newLoop.getResults().take_front(
                                                   loop.getNumResults())))

diff  --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index a6beca824d698..121386a09ad72 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -1279,9 +1279,10 @@ void AffineForOp::build(OpBuilder &builder, OperationState &result,
   Region *bodyRegion = result.addRegion();
   bodyRegion->push_back(new Block);
   Block &bodyBlock = bodyRegion->front();
-  Value inductionVar = bodyBlock.addArgument(builder.getIndexType());
+  Value inductionVar =
+      bodyBlock.addArgument(builder.getIndexType(), result.location);
   for (Value val : iterArgs)
-    bodyBlock.addArgument(val.getType());
+    bodyBlock.addArgument(val.getType(), val.getLoc());
 
   // Create the default terminator if the builder is not provided and if the
   // iteration arguments are not provided. Otherwise, leave this to the caller
@@ -1965,7 +1966,7 @@ AffineForOp mlir::replaceForOpWithNewYields(OpBuilder &b, AffineForOp loop,
   // Take the body of the original parent loop.
   newLoop.getLoopBody().takeBody(loop.getLoopBody());
   for (Value val : newIterArgs)
-    newLoop.getLoopBody().addArgument(val.getType());
+    newLoop.getLoopBody().addArgument(val.getType(), val.getLoc());
 
   // Update yield operation with new values to be added.
   if (!newYieldedValues.empty()) {
@@ -2887,7 +2888,7 @@ void AffineParallelOp::build(OpBuilder &builder, OperationState &result,
   auto *body = new Block();
   // Add all the block arguments.
   for (unsigned i = 0, e = steps.size(); i < e; ++i)
-    body->addArgument(IndexType::get(builder.getContext()));
+    body->addArgument(IndexType::get(builder.getContext()), result.location);
   bodyRegion->push_back(body);
   if (resultTypes.empty())
     ensureTerminator(*bodyRegion, builder, result.location);

diff  --git a/mlir/lib/Dialect/Async/IR/Async.cpp b/mlir/lib/Dialect/Async/IR/Async.cpp
index b7b708b7a8cb3..86036749eb64f 100644
--- a/mlir/lib/Dialect/Async/IR/Async.cpp
+++ b/mlir/lib/Dialect/Async/IR/Async.cpp
@@ -107,7 +107,8 @@ void ExecuteOp::build(OpBuilder &builder, OperationState &result,
   for (Value operand : operands) {
     auto valueType = operand.getType().dyn_cast<ValueType>();
     bodyBlock.addArgument(valueType ? valueType.getValueType()
-                                    : operand.getType());
+                                    : operand.getType(),
+                          operand.getLoc());
   }
 
   // Create the default terminator if the builder is not provided and if the

diff  --git a/mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp b/mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp
index db09c11574acc..edcbefc8e977a 100644
--- a/mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp
+++ b/mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp
@@ -270,7 +270,9 @@ static ParallelComputeFunction createParallelComputeFunction(
   rewriter.getListener()->notifyOperationInserted(func);
 
   // Create function entry block.
-  Block *block = b.createBlock(&func.getBody(), func.begin(), type.getInputs());
+  Block *block =
+      b.createBlock(&func.getBody(), func.begin(), type.getInputs(),
+                    SmallVector<Location>(type.getNumInputs(), op.getLoc()));
   b.setInsertionPointToEnd(block);
 
   ParallelComputeFunctionArgs args = {op.getNumLoops(), func.getArguments()};
@@ -482,7 +484,8 @@ static FuncOp createAsyncDispatchFunction(ParallelComputeFunction &computeFunc,
   rewriter.getListener()->notifyOperationInserted(func);
 
   // Create function entry block.
-  Block *block = b.createBlock(&func.getBody(), func.begin(), type.getInputs());
+  Block *block = b.createBlock(&func.getBody(), func.begin(), type.getInputs(),
+                               SmallVector<Location>(type.getNumInputs(), loc));
   b.setInsertionPointToEnd(block);
 
   Type indexTy = b.getIndexType();
@@ -499,11 +502,12 @@ static FuncOp createAsyncDispatchFunction(ParallelComputeFunction &computeFunc,
   // Create a work splitting while loop for the [blockStart, blockEnd) range.
   SmallVector<Type> types = {indexTy, indexTy};
   SmallVector<Value> operands = {blockStart, blockEnd};
+  SmallVector<Location> locations = {loc, loc};
 
   // Create a recursive dispatch loop.
   scf::WhileOp whileOp = b.create<scf::WhileOp>(types, operands);
-  Block *before = b.createBlock(&whileOp.getBefore(), {}, types);
-  Block *after = b.createBlock(&whileOp.getAfter(), {}, types);
+  Block *before = b.createBlock(&whileOp.getBefore(), {}, types, locations);
+  Block *after = b.createBlock(&whileOp.getAfter(), {}, types, locations);
 
   // Setup dispatch loop condition block: decide if we need to go into the
   // `after` block and launch one more async dispatch.

diff  --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index 429b7ffe1168c..75017b9a50d52 100644
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -369,8 +369,8 @@ void LaunchOp::build(OpBuilder &builder, OperationState &result,
   // the rest have the same types as the data operands.
   Region *kernelRegion = result.addRegion();
   Block *body = new Block();
-  body->addArguments(
-      std::vector<Type>(kNumConfigRegionAttributes, builder.getIndexType()));
+  for (unsigned i = 0; i < kNumConfigRegionAttributes; ++i)
+    body->addArgument(builder.getIndexType(), result.location);
   kernelRegion->push_back(body);
 }
 
@@ -671,7 +671,7 @@ parseLaunchFuncOperands(OpAsmParser &parser,
   if (parser.parseOptionalKeyword("args"))
     return success();
   SmallVector<NamedAttrList> argAttrs;
-  SmallVector<Optional<Location>> argLocations;
+  SmallVector<Location> argLocations;
   bool isVariadic = false;
   return function_interface_impl::parseFunctionArgumentList(
       parser, /*allowAttributes=*/false,
@@ -699,21 +699,21 @@ static void printLaunchFuncOperands(OpAsmPrinter &printer, Operation *,
 
 /// Adds a new block argument that corresponds to buffers located in
 /// workgroup memory.
-BlockArgument GPUFuncOp::addWorkgroupAttribution(Type type) {
+BlockArgument GPUFuncOp::addWorkgroupAttribution(Type type, Location loc) {
   auto attrName = getNumWorkgroupAttributionsAttrName();
   auto attr = (*this)->getAttrOfType<IntegerAttr>(attrName);
   (*this)->setAttr(attrName,
                    IntegerAttr::get(attr.getType(), attr.getValue() + 1));
   return getBody().insertArgument(getType().getNumInputs() + attr.getInt(),
-                                  type);
+                                  type, loc);
 }
 
 /// Adds a new block argument that corresponds to buffers located in
 /// private memory.
-BlockArgument GPUFuncOp::addPrivateAttribution(Type type) {
+BlockArgument GPUFuncOp::addPrivateAttribution(Type type, Location loc) {
   // Buffers on the private memory always come after buffers on the workgroup
   // memory.
-  return getBody().addArgument(type);
+  return getBody().addArgument(type, loc);
 }
 
 void GPUFuncOp::build(OpBuilder &builder, OperationState &result,
@@ -729,9 +729,14 @@ void GPUFuncOp::build(OpBuilder &builder, OperationState &result,
   result.addAttributes(attrs);
   Region *body = result.addRegion();
   Block *entryBlock = new Block;
-  entryBlock->addArguments(type.getInputs());
-  entryBlock->addArguments(workgroupAttributions);
-  entryBlock->addArguments(privateAttributions);
+
+  // TODO: Allow passing in proper locations here.
+  for (Type type : type.getInputs())
+    entryBlock->addArgument(type, result.location);
+  for (Type type : workgroupAttributions)
+    entryBlock->addArgument(type, result.location);
+  for (Type type : privateAttributions)
+    entryBlock->addArgument(type, result.location);
 
   body->getBlocks().push_back(entryBlock);
 }
@@ -783,7 +788,7 @@ static ParseResult parseGPUFuncOp(OpAsmParser &parser, OperationState &result) {
   SmallVector<NamedAttrList> resultAttrs;
   SmallVector<Type> argTypes;
   SmallVector<Type> resultTypes;
-  SmallVector<Optional<Location>> argLocations;
+  SmallVector<Location> argLocations;
   bool isVariadic;
 
   // Parse the function name.

diff  --git a/mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp b/mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp
index 7757e654944ca..81c8772cc2063 100644
--- a/mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/AllReduceLowering.cpp
@@ -157,11 +157,12 @@ struct GpuAllReduceRewriter {
 
   /// Adds type to funcOp's workgroup attributions.
   Value createWorkgroupBuffer() {
+    // TODO: Pick a proper location for the attribution.
     int workgroupMemoryAddressSpace =
         gpu::GPUDialect::getWorkgroupAddressSpace();
     auto bufferType = MemRefType::get({kSubgroupSize}, valueType, AffineMap{},
                                       workgroupMemoryAddressSpace);
-    return funcOp.addWorkgroupAttribution(bufferType);
+    return funcOp.addWorkgroupAttribution(bufferType, rewriter.getUnknownLoc());
   }
 
   /// Returns an accumulator factory using either the op attribute or the body
@@ -207,7 +208,7 @@ struct GpuAllReduceRewriter {
 
       // Return accumulator result.
       rewriter.setInsertionPointToStart(split);
-      return split->addArgument(lhs.getType());
+      return split->addArgument(lhs.getType(), lhs.getLoc());
     });
   }
 
@@ -298,7 +299,7 @@ struct GpuAllReduceRewriter {
     assert(thenOperands.size() == elseOperands.size());
     rewriter.setInsertionPointToStart(continueBlock);
     for (auto operand : thenOperands)
-      continueBlock->addArgument(operand.getType());
+      continueBlock->addArgument(operand.getType(), operand.getLoc());
   }
 
   /// Shortcut for createIf with empty else block and no block operands.

diff  --git a/mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp b/mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp
index fc00b5fc67cb6..5d7ded743da59 100644
--- a/mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp
@@ -252,7 +252,9 @@ struct GpuAsyncRegionPass::DeferWaitCallback {
           executeOp.operandsMutable().append(asyncTokens);
           SmallVector<Type, 1> tokenTypes(
               asyncTokens.size(), builder.getType<gpu::AsyncTokenType>());
-          copy(executeOp.getBody()->addArguments(tokenTypes),
+          SmallVector<Location, 1> tokenLocs(asyncTokens.size(),
+                                             executeOp.getLoc());
+          copy(executeOp.getBody()->addArguments(tokenTypes, tokenLocs),
                std::back_inserter(tokens));
         });
 

diff  --git a/mlir/lib/Dialect/GPU/Transforms/MemoryPromotion.cpp b/mlir/lib/Dialect/GPU/Transforms/MemoryPromotion.cpp
index 1e33c952ecaa5..182a24cd04444 100644
--- a/mlir/lib/Dialect/GPU/Transforms/MemoryPromotion.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/MemoryPromotion.cpp
@@ -150,8 +150,7 @@ void mlir::promoteToWorkgroupMemory(GPUFuncOp op, unsigned arg) {
   int workgroupMemoryAddressSpace = gpu::GPUDialect::getWorkgroupAddressSpace();
   auto bufferType = MemRefType::get(type.getShape(), type.getElementType(), {},
                                     workgroupMemoryAddressSpace);
-
-  Value attribution = op.addWorkgroupAttribution(bufferType);
+  Value attribution = op.addWorkgroupAttribution(bufferType, value.getLoc());
 
   // Replace the uses first since only the original uses are currently present.
   // Then insert the copies.

diff  --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index 2d0093970e8fd..440fc835e7cbb 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -1952,9 +1952,10 @@ Block *LLVMFuncOp::addEntryBlock() {
   auto *entry = new Block;
   push_back(entry);
 
+  // FIXME: Allow passing in proper locations for the entry arguments.
   LLVMFunctionType type = getType();
   for (unsigned i = 0, e = type.getNumParams(); i < e; ++i)
-    entry->addArgument(type.getParamType(i));
+    entry->addArgument(type.getParamType(i), getLoc());
   return entry;
 }
 
@@ -2038,7 +2039,7 @@ static ParseResult parseLLVMFuncOp(OpAsmParser &parser,
   SmallVector<NamedAttrList> resultAttrs;
   SmallVector<Type> argTypes;
   SmallVector<Type> resultTypes;
-  SmallVector<Optional<Location>> argLocations;
+  SmallVector<Location> argLocations;
   bool isVariadic;
 
   auto signatureLocation = parser.getCurrentLocation();

diff  --git a/mlir/lib/Dialect/LLVMIR/Transforms/LegalizeForExport.cpp b/mlir/lib/Dialect/LLVMIR/Transforms/LegalizeForExport.cpp
index f4fa100e6a78f..08fd412ff92cf 100644
--- a/mlir/lib/Dialect/LLVMIR/Transforms/LegalizeForExport.cpp
+++ b/mlir/lib/Dialect/LLVMIR/Transforms/LegalizeForExport.cpp
@@ -48,7 +48,8 @@ static void ensureDistinctSuccessors(Block &bb) {
     for (int position : llvm::drop_begin(successor.second, 1)) {
       Block *dummyBlock = builder.createBlock(bb.getParent());
       terminator->setSuccessor(dummyBlock, position);
-      dummyBlock->addArguments(successor.first->getArgumentTypes());
+      for (BlockArgument arg : successor.first->getArguments())
+        dummyBlock->addArgument(arg.getType(), arg.getLoc());
       builder.create<LLVM::BrOp>(terminator->getLoc(),
                                  dummyBlock->getArguments(), successor.first);
     }

diff  --git a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp
index fae2682c239fb..3c13e0a748f23 100644
--- a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp
+++ b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp
@@ -444,7 +444,7 @@ static LogicalResult bufferizeFuncOpBoundary(FuncOp funcOp,
     auto tensorType = bbArg.getType().dyn_cast<TensorType>();
     // Non-tensor types are just forwarded.
     if (!tensorType) {
-      frontBlock.addArgument(bbArg.getType());
+      frontBlock.addArgument(bbArg.getType(), bbArg.getLoc());
       bbArg.replaceAllUsesWith(frontBlock.getArguments().back());
       frontBlock.eraseArgument(0);
       continue;
@@ -452,7 +452,7 @@ static LogicalResult bufferizeFuncOpBoundary(FuncOp funcOp,
 
     // Get the buffer type from the bufferized function type.
     Type memrefType = bufferizedFuncType.getInput(idx);
-    Value memref = frontBlock.addArgument(memrefType);
+    Value memref = frontBlock.addArgument(memrefType, bbArg.getLoc());
     OpBuilder b(funcOp->getContext());
     b.setInsertionPointToStart(&frontBlock);
     // Replace all uses of bbArg through a ToMemRefOp by a memref::CastOp.

diff  --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 00a120fb45189..f41e8487858f3 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -128,8 +128,8 @@ static LogicalResult foldMemRefCastInTiledLoopOp(TiledLoopOp op) {
     auto castOp = operand.get().getDefiningOp<memref::CastOp>();
     if (castOp && memref::CastOp::canFoldIntoConsumerOp(castOp)) {
       operand.set(castOp.getOperand());
-      BlockArgument newBbArg =
-          body->insertArgument(bbArgIndex, castOp.getOperand().getType());
+      BlockArgument newBbArg = body->insertArgument(
+          bbArgIndex, castOp.getOperand().getType(), op.getLoc());
       BlockArgument oldBbArg = body->getArgument(newBbArg.getArgNumber() + 1);
 
       // Insert memref.cast back to the original type.
@@ -573,13 +573,18 @@ void GenericOp::build(
     return;
 
   SmallVector<Type, 4> blockArgTypes;
-  for (ValueRange container : {inputs, outputs})
-    for (Value v : container)
+  SmallVector<Location, 4> blockArgLocs;
+  for (ValueRange container : {inputs, outputs}) {
+    for (Value v : container) {
       blockArgTypes.push_back(getElementTypeOrSelf(v));
+      blockArgLocs.push_back(v.getLoc());
+    }
+  }
 
   OpBuilder::InsertionGuard guard(builder);
   auto &region = *result.regions.front();
-  Block *bodyBlock = builder.createBlock(&region, region.end(), blockArgTypes);
+  Block *bodyBlock =
+      builder.createBlock(&region, region.end(), blockArgTypes, blockArgLocs);
   bodyBuild(builder, result.location, bodyBlock->getArguments());
 }
 
@@ -1208,13 +1213,13 @@ PadTensorOp PadTensorOp::createPadScalarOp(Type type, Value source, Value pad,
   auto padTensorOp =
       builder.create<linalg::PadTensorOp>(loc, type, source, low, high, nofold);
   int rank = padTensorOp.getResultType().getRank();
-  SmallVector<Type, 4> blockArgTypes;
-  blockArgTypes.assign(rank, builder.getIndexType());
+  SmallVector<Type, 4> blockArgTypes(rank, builder.getIndexType());
+  SmallVector<Location, 4> blockArgLocs(rank, loc);
   auto &region = padTensorOp.region();
   // `builder.createBlock` changes the insertion point within the block. Create
   // a guard to reset the insertion point of the builder after it is destroyed.
   OpBuilder::InsertionGuard guard(builder);
-  builder.createBlock(&region, region.end(), blockArgTypes);
+  builder.createBlock(&region, region.end(), blockArgTypes, blockArgLocs);
   builder.create<linalg::YieldOp>(loc, pad);
   return padTensorOp;
 }
@@ -1770,12 +1775,17 @@ void TiledLoopOp::build(OpBuilder &builder, OperationState &result,
   OpBuilder::InsertionGuard guard(builder);
   unsigned numIVs = steps.size();
   SmallVector<Type, 8> argTypes(numIVs, builder.getIndexType());
-  for (Type type : TypeRange(inputs))
-    argTypes.push_back(type);
-  for (Type type : TypeRange(outputs))
-    argTypes.push_back(type);
+  SmallVector<Location, 8> argLocs(numIVs, result.location);
+  for (Value input : inputs) {
+    argTypes.push_back(input.getType());
+    argLocs.push_back(input.getLoc());
+  }
+  for (Value output : outputs) {
+    argTypes.push_back(output.getType());
+    argLocs.push_back(output.getLoc());
+  }
   Region *bodyRegion = result.addRegion();
-  Block *bodyBlock = builder.createBlock(bodyRegion, {}, argTypes);
+  Block *bodyBlock = builder.createBlock(bodyRegion, {}, argTypes, argLocs);
 
   if (bodyBuilderFn) {
     builder.setInsertionPointToStart(bodyBlock);
@@ -2437,13 +2447,20 @@ static void fillStructuredOpRegion(
   // TODO: atm all operands go through getElementTypeOrSelf,
   // reconsider when we have evidence we need to.
   SmallVector<Type, 8> argTypes;
-  for (auto containers : {inputTypes, outputTypes})
-    for (auto t : containers)
+  SmallVector<Location, 8> argLocs;
+  for (auto containers : {inputTypes, outputTypes}) {
+    for (auto t : containers) {
       argTypes.push_back(getElementTypeOrSelf(t));
 
+      // TODO: Pass in a proper location here.
+      argLocs.push_back(opBuilder.getUnknownLoc());
+    }
+  }
+
   // RAII.
   OpBuilder::InsertionGuard guard(opBuilder);
-  Block *body = opBuilder.createBlock(&region, /*insertPt=*/{}, argTypes);
+  Block *body =
+      opBuilder.createBlock(&region, /*insertPt=*/{}, argTypes, argLocs);
   unsigned actual = body->getNumArguments();
   unsigned expected = NamedStructuredOpType::getNumRegionArgs();
   if (expected != actual) {

diff  --git a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
index 619908ca5a385..33286258543e5 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
@@ -172,7 +172,7 @@ generateFusedElementwiseOpRegion(PatternRewriter &rewriter, GenericOp fusedOp,
   // 3. Consumer input operands up to consumerIdx (exclusive).
   for (BlockArgument bbArg : consumerBlock.getArguments().take_front(
            consumerOpOperand->getOperandNumber())) // input assumption.
-    mapper.map(bbArg, fusedBlock->addArgument(bbArg.getType()));
+    mapper.map(bbArg, fusedBlock->addArgument(bbArg.getType(), bbArg.getLoc()));
 
   // Replacing consumerIdx requires getting the cloned, yielded, value from
   // the (cloned) producer block. This happens in step 9.
@@ -180,7 +180,7 @@ generateFusedElementwiseOpRegion(PatternRewriter &rewriter, GenericOp fusedOp,
   // 4. Splice in producer's input operands.
   for (BlockArgument bbArg :
        producerBlock.getArguments().take_front(producer.getNumInputs()))
-    mapper.map(bbArg, fusedBlock->addArgument(bbArg.getType()));
+    mapper.map(bbArg, fusedBlock->addArgument(bbArg.getType(), bbArg.getLoc()));
 
   // 4.b. Producer output operand/map that is fused needs to be mapped to the
   // producer bbArg if it is an "initTensor" (i.e. its value is actually read).
@@ -190,18 +190,18 @@ generateFusedElementwiseOpRegion(PatternRewriter &rewriter, GenericOp fusedOp,
                               .drop_front(producer.getNumInputs())
                               // TODO: bbArg index of
                               .front();
-    mapper.map(bbArg, fusedBlock->addArgument(bbArg.getType()));
+    mapper.map(bbArg, fusedBlock->addArgument(bbArg.getType(), bbArg.getLoc()));
   }
   // 5. Remaining consumer's input operands (drop past index `consumerIdx`).
   for (BlockArgument bbArg :
        consumerBlock.getArguments()
            .take_front(consumer.getNumInputs())
            .drop_front(consumerOpOperand->getOperandNumber() + 1))
-    mapper.map(bbArg, fusedBlock->addArgument(bbArg.getType()));
+    mapper.map(bbArg, fusedBlock->addArgument(bbArg.getType(), bbArg.getLoc()));
   // 6. All of consumer's output operands.
   for (BlockArgument bbArg :
        consumerBlock.getArguments().take_back(consumer.getNumOutputs()))
-    mapper.map(bbArg, fusedBlock->addArgument(bbArg.getType()));
+    mapper.map(bbArg, fusedBlock->addArgument(bbArg.getType(), bbArg.getLoc()));
   // 7. All of producer's output operands except the one fused.
   // TODO: allow fusion of multi-result producers.
   assert(producer->getNumResults() == 1 && "expected single result producer");

diff  --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 553775cd7faf2..a527954dfd8ec 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -1291,7 +1291,8 @@ void WsLoopOp::build(OpBuilder &builder, OperationState &result,
     OpBuilder::InsertionGuard guard(builder);
     unsigned numIVs = steps.size();
     SmallVector<Type, 8> argTypes(numIVs, steps.getType().front());
-    builder.createBlock(bodyRegion, {}, argTypes);
+    SmallVector<Location, 8> argLocs(numIVs, result.location);
+    builder.createBlock(bodyRegion, {}, argTypes, argLocs);
   }
 }
 

diff  --git a/mlir/lib/Dialect/PDLInterp/IR/PDLInterp.cpp b/mlir/lib/Dialect/PDLInterp/IR/PDLInterp.cpp
index c7b67e859c8ac..20d6a376b47d2 100644
--- a/mlir/lib/Dialect/PDLInterp/IR/PDLInterp.cpp
+++ b/mlir/lib/Dialect/PDLInterp/IR/PDLInterp.cpp
@@ -74,9 +74,11 @@ void ForEachOp::build(::mlir::OpBuilder &builder, ::mlir::OperationState &state,
   build(builder, state, range, successor);
   if (initLoop) {
     // Create the block and the loop variable.
+    // FIXME: Allow passing in a proper location for the loop variable.
     auto rangeType = range.getType().cast<pdl::RangeType>();
     state.regions.front()->emplaceBlock();
-    state.regions.front()->addArgument(rangeType.getElementType());
+    state.regions.front()->addArgument(rangeType.getElementType(),
+                                       state.location);
   }
 }
 

diff  --git a/mlir/lib/Dialect/SCF/SCF.cpp b/mlir/lib/Dialect/SCF/SCF.cpp
index ac70c57fea5e6..7276ed0cdb740 100644
--- a/mlir/lib/Dialect/SCF/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/SCF.cpp
@@ -224,7 +224,7 @@ struct MultiBlockExecuteInliner : public OpRewritePattern<ExecuteRegionOp> {
     SmallVector<Value> blockArgs;
 
     for (auto res : op.getResults())
-      blockArgs.push_back(postBlock->addArgument(res.getType()));
+      blockArgs.push_back(postBlock->addArgument(res.getType(), res.getLoc()));
 
     rewriter.replaceOp(op, blockArgs);
     return success();
@@ -260,9 +260,9 @@ void ForOp::build(OpBuilder &builder, OperationState &result, Value lb,
   Region *bodyRegion = result.addRegion();
   bodyRegion->push_back(new Block);
   Block &bodyBlock = bodyRegion->front();
-  bodyBlock.addArgument(builder.getIndexType());
+  bodyBlock.addArgument(builder.getIndexType(), result.location);
   for (Value v : iterArgs)
-    bodyBlock.addArgument(v.getType());
+    bodyBlock.addArgument(v.getType(), v.getLoc());
 
   // Create the default terminator if the builder is not provided and if the
   // iteration arguments are not provided. Otherwise, leave this to the caller
@@ -1677,8 +1677,9 @@ void ParallelOp::build(
   OpBuilder::InsertionGuard guard(builder);
   unsigned numIVs = steps.size();
   SmallVector<Type, 8> argTypes(numIVs, builder.getIndexType());
+  SmallVector<Location, 8> argLocs(numIVs, result.location);
   Region *bodyRegion = result.addRegion();
-  Block *bodyBlock = builder.createBlock(bodyRegion, {}, argTypes);
+  Block *bodyBlock = builder.createBlock(bodyRegion, {}, argTypes, argLocs);
 
   if (bodyBuilderFn) {
     builder.setInsertionPointToStart(bodyBlock);
@@ -2054,7 +2055,8 @@ void ReduceOp::build(
 
   OpBuilder::InsertionGuard guard(builder);
   Region *bodyRegion = result.addRegion();
-  Block *body = builder.createBlock(bodyRegion, {}, ArrayRef<Type>{type, type});
+  Block *body = builder.createBlock(bodyRegion, {}, ArrayRef<Type>{type, type},
+                                    {result.location, result.location});
   if (bodyBuilderFn)
     bodyBuilderFn(builder, result.location, body->getArgument(0),
                   body->getArgument(1));
@@ -2368,6 +2370,7 @@ struct WhileUnusedResult : public OpRewritePattern<WhileOp> {
     SmallVector<unsigned> newResultsIndices;
     SmallVector<Type> newResultTypes;
     SmallVector<Value> newTermArgs;
+    SmallVector<Location> newArgLocs;
     bool needUpdate = false;
     for (const auto &it :
          llvm::enumerate(llvm::zip(op.getResults(), afterArgs, termArgs))) {
@@ -2381,6 +2384,7 @@ struct WhileUnusedResult : public OpRewritePattern<WhileOp> {
         newResultsIndices.emplace_back(i);
         newTermArgs.emplace_back(termArg);
         newResultTypes.emplace_back(result.getType());
+        newArgLocs.emplace_back(result.getLoc());
       }
     }
 
@@ -2398,7 +2402,7 @@ struct WhileUnusedResult : public OpRewritePattern<WhileOp> {
         rewriter.create<WhileOp>(op.getLoc(), newResultTypes, op.getInits());
 
     Block &newAfterBlock = *rewriter.createBlock(
-        &newWhile.getAfter(), /*insertPt*/ {}, newResultTypes);
+        &newWhile.getAfter(), /*insertPt*/ {}, newResultTypes, newArgLocs);
 
     // Build new results list and new after block args (unused entries will be
     // null).

diff  --git a/mlir/lib/Dialect/SCF/Transforms/ForToWhile.cpp b/mlir/lib/Dialect/SCF/Transforms/ForToWhile.cpp
index a297c7bb06b4a..faf0da11302c8 100644
--- a/mlir/lib/Dialect/SCF/Transforms/ForToWhile.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ForToWhile.cpp
@@ -33,10 +33,14 @@ struct ForLoopLoweringPattern : public OpRewritePattern<ForOp> {
                                 PatternRewriter &rewriter) const override {
     // Generate type signature for the loop-carried values. The induction
     // variable is placed first, followed by the forOp.iterArgs.
-    SmallVector<Type, 8> lcvTypes;
+    SmallVector<Type> lcvTypes;
+    SmallVector<Location> lcvLocs;
     lcvTypes.push_back(forOp.getInductionVar().getType());
-    llvm::transform(forOp.getInitArgs(), std::back_inserter(lcvTypes),
-                    [&](auto v) { return v.getType(); });
+    lcvLocs.push_back(forOp.getInductionVar().getLoc());
+    for (Value value : forOp.getInitArgs()) {
+      lcvTypes.push_back(value.getType());
+      lcvLocs.push_back(value.getLoc());
+    }
 
     // Build scf.WhileOp
     SmallVector<Value> initArgs;
@@ -48,7 +52,7 @@ struct ForLoopLoweringPattern : public OpRewritePattern<ForOp> {
     // 'before' region contains the loop condition and forwarding of iteration
     // arguments to the 'after' region.
     auto *beforeBlock = rewriter.createBlock(
-        &whileOp.getBefore(), whileOp.getBefore().begin(), lcvTypes, {});
+        &whileOp.getBefore(), whileOp.getBefore().begin(), lcvTypes, lcvLocs);
     rewriter.setInsertionPointToStart(&whileOp.getBefore().front());
     auto cmpOp = rewriter.create<arith::CmpIOp>(
         whileOp.getLoc(), arith::CmpIPredicate::slt,
@@ -60,7 +64,7 @@ struct ForLoopLoweringPattern : public OpRewritePattern<ForOp> {
     // region. The return type of the execRegionOp does not contain the
     // iv - yields in the source for-loop contain only iterArgs.
     auto *afterBlock = rewriter.createBlock(
-        &whileOp.getAfter(), whileOp.getAfter().begin(), lcvTypes, {});
+        &whileOp.getAfter(), whileOp.getAfter().begin(), lcvTypes, lcvLocs);
 
     // Add induction variable incrementation
     rewriter.setInsertionPointToEnd(afterBlock);

diff  --git a/mlir/lib/Dialect/SCF/Transforms/Utils.cpp b/mlir/lib/Dialect/SCF/Transforms/Utils.cpp
index abd7408f86cd1..66be70bf1a19d 100644
--- a/mlir/lib/Dialect/SCF/Transforms/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/Utils.cpp
@@ -106,11 +106,18 @@ FailureOr<FuncOp> mlir::outlineSingleBlockRegion(RewriterBase &rewriter,
 
   ValueRange outlinedValues(captures.getArrayRef());
   SmallVector<Type> outlinedFuncArgTypes;
+  SmallVector<Location> outlinedFuncArgLocs;
   // Region's arguments are exactly the first block's arguments as per
   // Region::getArguments().
   // Func's arguments are cat(regions's arguments, captures arguments).
-  llvm::append_range(outlinedFuncArgTypes, region.getArgumentTypes());
-  llvm::append_range(outlinedFuncArgTypes, outlinedValues.getTypes());
+  for (BlockArgument arg : region.getArguments()) {
+    outlinedFuncArgTypes.push_back(arg.getType());
+    outlinedFuncArgLocs.push_back(arg.getLoc());
+  }
+  for (Value value : outlinedValues) {
+    outlinedFuncArgTypes.push_back(value.getType());
+    outlinedFuncArgLocs.push_back(value.getLoc());
+  }
   FunctionType outlinedFuncType =
       FunctionType::get(rewriter.getContext(), outlinedFuncArgTypes,
                         originalTerminator->getOperandTypes());
@@ -137,7 +144,9 @@ FailureOr<FuncOp> mlir::outlineSingleBlockRegion(RewriterBase &rewriter,
   // terminator(call_results).
   Block *newBlock = rewriter.createBlock(
       &region, region.begin(),
-      TypeRange{outlinedFuncArgTypes}.take_front(numOriginalBlockArguments));
+      TypeRange{outlinedFuncArgTypes}.take_front(numOriginalBlockArguments),
+      ArrayRef<Location>(outlinedFuncArgLocs)
+          .take_front(numOriginalBlockArguments));
   {
     OpBuilder::InsertionGuard g(rewriter);
     rewriter.setInsertionPointToEnd(newBlock);

diff  --git a/mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp b/mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp
index ca4ec92b62e9f..cf764eaa71c47 100644
--- a/mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp
@@ -1947,7 +1947,7 @@ static ParseResult parseFuncOp(OpAsmParser &parser, OperationState &state) {
   SmallVector<NamedAttrList> resultAttrs;
   SmallVector<Type> argTypes;
   SmallVector<Type> resultTypes;
-  SmallVector<Optional<Location>> argLocations;
+  SmallVector<Location> argLocations;
   auto &builder = parser.getBuilder();
 
   // Parse the name as a symbol.

diff  --git a/mlir/lib/Dialect/Shape/IR/Shape.cpp b/mlir/lib/Dialect/Shape/IR/Shape.cpp
index d1263a1586af0..7e97846c33572 100644
--- a/mlir/lib/Dialect/Shape/IR/Shape.cpp
+++ b/mlir/lib/Dialect/Shape/IR/Shape.cpp
@@ -1625,18 +1625,18 @@ void ReduceOp::build(OpBuilder &builder, OperationState &result, Value shape,
   Region *bodyRegion = result.addRegion();
   bodyRegion->push_back(new Block);
   Block &bodyBlock = bodyRegion->front();
-  bodyBlock.addArgument(builder.getIndexType());
+  bodyBlock.addArgument(builder.getIndexType(), result.location);
 
   Type elementType;
   if (auto tensorType = shape.getType().dyn_cast<TensorType>())
     elementType = tensorType.getElementType();
   else
     elementType = SizeType::get(builder.getContext());
-  bodyBlock.addArgument(elementType);
+  bodyBlock.addArgument(elementType, shape.getLoc());
 
-  for (Type initValType : initVals.getTypes()) {
-    bodyBlock.addArgument(initValType);
-    result.addTypes(initValType);
+  for (Value initVal : initVals) {
+    bodyBlock.addArgument(initVal.getType(), initVal.getLoc());
+    result.addTypes(initVal.getType());
   }
 }
 

diff  --git a/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
index 9afb67ef0ceef..8284dd97b3d25 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
@@ -1204,8 +1204,10 @@ static Operation *genWhile(Merger &merger, CodeGen &codegen,
   assert(types.size() == operands.size());
   Location loc = op.getLoc();
   scf::WhileOp whileOp = rewriter.create<scf::WhileOp>(loc, types, operands);
-  Block *before = rewriter.createBlock(&whileOp.getBefore(), {}, types);
-  Block *after = rewriter.createBlock(&whileOp.getAfter(), {}, types);
+
+  SmallVector<Location> locs(types.size(), loc);
+  Block *before = rewriter.createBlock(&whileOp.getBefore(), {}, types, locs);
+  Block *after = rewriter.createBlock(&whileOp.getAfter(), {}, types, locs);
 
   // Build the "before" region, which effectively consists
   // of a conjunction of "i < upper" tests on all induction.

diff  --git a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
index 656394009cbc4..1110879395259 100644
--- a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
+++ b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
@@ -146,7 +146,7 @@ void GenericAtomicRMWOp::build(OpBuilder &builder, OperationState &result,
 
     Region *bodyRegion = result.addRegion();
     bodyRegion->push_back(new Block());
-    bodyRegion->addArgument(elementType);
+    bodyRegion->addArgument(elementType, memref.getLoc());
   }
 }
 

diff  --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index cf3122c8bcd70..613edde638683 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -495,8 +495,9 @@ void GenerateOp::build(
   Region *bodyRegion = result.regions.front().get();
   auto rank = resultTy.cast<RankedTensorType>().getRank();
   SmallVector<Type, 2> argumentTypes(rank, b.getIndexType());
+  SmallVector<Location, 2> argumentLocs(rank, result.location);
   Block *bodyBlock =
-      b.createBlock(bodyRegion, bodyRegion->end(), argumentTypes);
+      b.createBlock(bodyRegion, bodyRegion->end(), argumentTypes, argumentLocs);
   bodyBuilder(b, result.location, bodyBlock->getArguments());
 }
 

diff  --git a/mlir/lib/IR/Block.cpp b/mlir/lib/IR/Block.cpp
index a138ec5d9cef0..172e8e4997365 100644
--- a/mlir/lib/IR/Block.cpp
+++ b/mlir/lib/IR/Block.cpp
@@ -138,17 +138,8 @@ auto Block::getArgumentTypes() -> ValueTypeRange<BlockArgListType> {
   return ValueTypeRange<BlockArgListType>(getArguments());
 }
 
-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);
+BlockArgument Block::addArgument(Type type, Location loc) {
+  BlockArgument arg = BlockArgument::create(type, this, arguments.size(), loc);
   arguments.push_back(arg);
   return arg;
 }
@@ -156,37 +147,20 @@ BlockArgument Block::addArgument(Type type, Optional<Location> loc) {
 /// Add one argument to the argument list for each type specified in the list.
 auto Block::addArguments(TypeRange types, ArrayRef<Location> locs)
     -> iterator_range<args_iterator> {
-  // TODO: Require locations for BlockArguments.
-  assert((locs.empty() || types.size() == locs.size()) &&
+  assert(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));
-  }
+  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,
-                                    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 Block::insertArgument(unsigned index, Type type, Location loc) {
+  assert(index <= arguments.size() && "invalid insertion index");
 
-  auto arg = BlockArgument::create(type, this, index, *loc);
-  assert(index <= arguments.size());
+  auto arg = BlockArgument::create(type, this, index, loc);
   arguments.insert(arguments.begin() + index, arg);
   // Update the cached position for all the arguments after the newly inserted
   // one.
@@ -198,8 +172,7 @@ BlockArgument Block::insertArgument(unsigned index, Type type,
 
 /// 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,
-                                    Optional<Location> loc) {
+BlockArgument Block::insertArgument(args_iterator it, Type type, Location loc) {
   assert(llvm::empty(getPredecessors()) &&
          "cannot insert arguments to blocks with predecessors");
   return insertArgument(it->getArgNumber(), type, loc);

diff  --git a/mlir/lib/IR/Builders.cpp b/mlir/lib/IR/Builders.cpp
index de7204f40979b..c49eceec8055f 100644
--- a/mlir/lib/IR/Builders.cpp
+++ b/mlir/lib/IR/Builders.cpp
@@ -350,12 +350,10 @@ Operation *OpBuilder::insert(Operation *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, ArrayRef<Location> locs) {
   assert(parent && "expected valid parent region");
+  assert(argTypes.size() == locs.size() && "argument location mismatch");
   if (insertPt == Region::iterator())
     insertPt = parent->end();
 

diff  --git a/mlir/lib/IR/FunctionImplementation.cpp b/mlir/lib/IR/FunctionImplementation.cpp
index 49ffaa231f217..dd3b40db781b2 100644
--- 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::OperandType> &argNames,
     SmallVectorImpl<Type> &argTypes, SmallVectorImpl<NamedAttrList> &argAttrs,
-    SmallVectorImpl<Optional<Location>> &argLocations, bool &isVariadic) {
+    SmallVectorImpl<Location> &argLocations, bool &isVariadic) {
   if (parser.parseLParen())
     return failure();
 
@@ -65,7 +65,9 @@ ParseResult mlir::function_interface_impl::parseFunctionArgumentList(
     if (!argument.name.empty() &&
         parser.parseOptionalLocationSpecifier(explicitLoc))
       return failure();
-    argLocations.push_back(explicitLoc);
+    if (!explicitLoc)
+      explicitLoc = parser.getEncodedSourceLoc(loc);
+    argLocations.push_back(*explicitLoc);
 
     return success();
   };
@@ -133,7 +135,7 @@ ParseResult mlir::function_interface_impl::parseFunctionSignature(
     OpAsmParser &parser, bool allowVariadic,
     SmallVectorImpl<OpAsmParser::OperandType> &argNames,
     SmallVectorImpl<Type> &argTypes, SmallVectorImpl<NamedAttrList> &argAttrs,
-    SmallVectorImpl<Optional<Location>> &argLocations, bool &isVariadic,
+    SmallVectorImpl<Location> &argLocations, bool &isVariadic,
     SmallVectorImpl<Type> &resultTypes,
     SmallVectorImpl<NamedAttrList> &resultAttrs) {
   bool allowArgAttrs = true;
@@ -197,7 +199,7 @@ ParseResult mlir::function_interface_impl::parseFunctionOp(
   SmallVector<NamedAttrList> resultAttrs;
   SmallVector<Type> argTypes;
   SmallVector<Type> resultTypes;
-  SmallVector<Optional<Location>> argLocations;
+  SmallVector<Location> argLocations;
   auto &builder = parser.getBuilder();
 
   // Parse visibility.
@@ -257,7 +259,7 @@ ParseResult mlir::function_interface_impl::parseFunctionOp(
   llvm::SMLoc loc = parser.getCurrentLocation();
   OptionalParseResult parseResult = parser.parseOptionalRegion(
       *body, entryArgs, entryArgs.empty() ? ArrayRef<Type>() : argTypes,
-      entryArgs.empty() ? ArrayRef<Optional<Location>>() : argLocations,
+      entryArgs.empty() ? ArrayRef<Location>() : argLocations,
       /*enableNameShadowing=*/false);
   if (parseResult.hasValue()) {
     if (failed(*parseResult))

diff  --git a/mlir/lib/IR/FunctionInterfaces.cpp b/mlir/lib/IR/FunctionInterfaces.cpp
index c7f3957cd3723..07da5ce1716f1 100644
--- a/mlir/lib/IR/FunctionInterfaces.cpp
+++ b/mlir/lib/IR/FunctionInterfaces.cpp
@@ -130,11 +130,11 @@ void mlir::function_interface_impl::setAllResultAttrDicts(
 
 void mlir::function_interface_impl::insertFunctionArguments(
     Operation *op, ArrayRef<unsigned> argIndices, TypeRange argTypes,
-    ArrayRef<DictionaryAttr> argAttrs, ArrayRef<Optional<Location>> argLocs,
+    ArrayRef<DictionaryAttr> argAttrs, ArrayRef<Location> argLocs,
     unsigned originalNumArgs, Type newType) {
   assert(argIndices.size() == argTypes.size());
   assert(argIndices.size() == argAttrs.size() || argAttrs.empty());
-  assert(argIndices.size() == argLocs.size() || argLocs.empty());
+  assert(argIndices.size() == argLocs.size());
   if (argIndices.empty())
     return;
 
@@ -171,8 +171,7 @@ void mlir::function_interface_impl::insertFunctionArguments(
   // Update the function type and any entry block arguments.
   op->setAttr(getTypeAttrName(), TypeAttr::get(newType));
   for (unsigned i = 0, e = argIndices.size(); i < e; ++i)
-    entry.insertArgument(argIndices[i] + i, argTypes[i],
-                         argLocs.empty() ? Optional<Location>{} : argLocs[i]);
+    entry.insertArgument(argIndices[i] + i, argTypes[i], argLocs[i]);
 }
 
 void mlir::function_interface_impl::insertFunctionResults(

diff  --git a/mlir/lib/IR/Region.cpp b/mlir/lib/IR/Region.cpp
index d1783c482d47f..8739b98d67018 100644
--- a/mlir/lib/IR/Region.cpp
+++ b/mlir/lib/IR/Region.cpp
@@ -33,14 +33,13 @@ Location Region::getLoc() {
   return container->getLoc();
 }
 
-/// Return a range containing the types of the arguments for this region.
 auto Region::getArgumentTypes() -> ValueTypeRange<BlockArgListType> {
   return ValueTypeRange<BlockArgListType>(getArguments());
 }
 
-/// Add one argument to the argument list for each type specified in the list.
-iterator_range<Region::args_iterator> Region::addArguments(TypeRange types) {
-  return front().addArguments(types);
+iterator_range<Region::args_iterator>
+Region::addArguments(TypeRange types, ArrayRef<Location> locs) {
+  return front().addArguments(types, locs);
 }
 
 Region *Region::getParentRegion() {

diff  --git a/mlir/lib/Parser/Parser.cpp b/mlir/lib/Parser/Parser.cpp
index 0cc0bbf471478..ad7abbb9368a7 100644
--- a/mlir/lib/Parser/Parser.cpp
+++ b/mlir/lib/Parser/Parser.cpp
@@ -368,15 +368,14 @@ class OperationParser : public Parser {
   /// region is isolated from those above.
   ParseResult parseRegion(Region &region,
                           ArrayRef<std::pair<SSAUseInfo, Type>> entryArguments,
-                          ArrayRef<Optional<Location>> argLocations = {},
+                          ArrayRef<Location> argLocations,
                           bool isIsolatedNameScope = false);
 
   /// Parse a region body into 'region'.
   ParseResult
   parseRegionBody(Region &region, llvm::SMLoc startLoc,
                   ArrayRef<std::pair<SSAUseInfo, Type>> entryArguments,
-                  ArrayRef<Optional<Location>> argLocations,
-                  bool isIsolatedNameScope);
+                  ArrayRef<Location> argLocations, bool isIsolatedNameScope);
 
   //===--------------------------------------------------------------------===//
   // Block Parsing
@@ -1055,7 +1054,8 @@ 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=*/{}))
+        if (parseRegion(*result.regions.back(), /*entryArguments=*/{},
+                        /*argLocations=*/{}))
           return failure();
       } while (consumeIf(Token::comma));
       if (parseToken(Token::r_paren, "expected ')' to end region list"))
@@ -1451,7 +1451,7 @@ class CustomOpAsmParser : public AsmParserImpl<OpAsmParser> {
   /// effectively defines the SSA values of `arguments` and assigns their type.
   ParseResult parseRegion(Region &region, ArrayRef<OperandType> arguments,
                           ArrayRef<Type> argTypes,
-                          ArrayRef<Optional<Location>> argLocations,
+                          ArrayRef<Location> argLocations,
                           bool enableNameShadowing) override {
     assert(arguments.size() == argTypes.size() &&
            "mismatching number of arguments and types");
@@ -1477,11 +1477,11 @@ class CustomOpAsmParser : public AsmParserImpl<OpAsmParser> {
   }
 
   /// Parses a region if present.
-  OptionalParseResult
-  parseOptionalRegion(Region &region, ArrayRef<OperandType> arguments,
-                      ArrayRef<Type> argTypes,
-                      ArrayRef<Optional<Location>> argLocations,
-                      bool enableNameShadowing) override {
+  OptionalParseResult parseOptionalRegion(Region &region,
+                                          ArrayRef<OperandType> 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,
@@ -1823,7 +1823,7 @@ OperationParser::parseTrailingLocationSpecifier(OpOrArgument opOrArgument) {
 ParseResult OperationParser::parseRegion(
     Region &region,
     ArrayRef<std::pair<OperationParser::SSAUseInfo, Type>> entryArguments,
-    ArrayRef<Optional<Location>> argLocations, bool isIsolatedNameScope) {
+    ArrayRef<Location> argLocations, bool isIsolatedNameScope) {
   // Parse the '{'.
   Token lBraceTok = getToken();
   if (parseToken(Token::l_brace, "expected '{' to begin a region"))
@@ -1851,7 +1851,7 @@ ParseResult OperationParser::parseRegion(
 ParseResult OperationParser::parseRegionBody(
     Region &region, llvm::SMLoc startLoc,
     ArrayRef<std::pair<OperationParser::SSAUseInfo, Type>> entryArguments,
-    ArrayRef<Optional<Location>> argLocations, bool isIsolatedNameScope) {
+    ArrayRef<Location> argLocations, bool isIsolatedNameScope) {
   assert(argLocations.empty() || argLocations.size() == entryArguments.size());
   auto currentPt = opBuilder.saveInsertionPoint();
 
@@ -1886,11 +1886,11 @@ ParseResult OperationParser::parseRegionBody(
                    .attachNote(getEncodedSourceLocation(*defLoc))
                << "previously referenced here";
       }
-      Location loc =
-          (!argLocations.empty() && argLocations[argIndex])
-              ? *argLocations[argIndex]
-              : getEncodedSourceLocation(placeholderArgPair.first.loc);
-      BlockArgument arg = block->addArgument(placeholderArgPair.second, loc);
+      BlockArgument arg = block->addArgument(
+          placeholderArgPair.second,
+          argLocations.empty()
+              ? getEncodedSourceLocation(placeholderArgPair.first.loc)
+              : argLocations[argIndex]);
 
       // Add a definition of this arg to the assembly state if provided.
       if (state.asmState)

diff  --git a/mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp b/mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp
index f7ad3383323b2..a41304200a307 100644
--- a/mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp
+++ b/mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp
@@ -647,7 +647,8 @@ LogicalResult Importer::processInstruction(llvm::Instruction *inst) {
     Type type = processType(inst->getType());
     if (!type)
       return failure();
-    v = b.getInsertionBlock()->addArgument(type);
+    v = b.getInsertionBlock()->addArgument(
+        type, processDebugLoc(inst->getDebugLoc(), inst));
     return success();
   }
   case llvm::Instruction::Call: {
@@ -822,9 +823,10 @@ LogicalResult Importer::processFunction(llvm::Function *f) {
   currentEntryBlock = blockList[0];
 
   // Add function arguments to the entry block.
-  for (const auto &kv : llvm::enumerate(f->args()))
-    instMap[&kv.value()] =
-        blockList[0]->addArgument(functionType.getParamType(kv.index()));
+  for (const auto &kv : llvm::enumerate(f->args())) {
+    instMap[&kv.value()] = blockList[0]->addArgument(
+        functionType.getParamType(kv.index()), fop.getLoc());
+  }
 
   for (auto bbs : llvm::zip(*f, blockList)) {
     if (failed(processBasicBlock(&std::get<0>(bbs), std::get<1>(bbs))))

diff  --git a/mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp b/mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp
index 8d055eed6a68c..361413ee24cfd 100644
--- a/mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp
+++ b/mlir/lib/Target/SPIRV/Deserialization/Deserializer.cpp
@@ -1580,7 +1580,7 @@ LogicalResult spirv::Deserializer::processPhi(ArrayRef<uint32_t> operands) {
 
   // Create a block argument for this OpPhi instruction.
   Type blockArgType = getType(operands[0]);
-  BlockArgument blockArg = curBlock->addArgument(blockArgType);
+  BlockArgument blockArg = curBlock->addArgument(blockArgType, unknownLoc);
   valueMap[operands[1]] = blockArg;
   LLVM_DEBUG(logger.startLine()
              << "[phi] created block argument " << blockArg
@@ -1748,7 +1748,8 @@ LogicalResult ControlFlowStructurizer::structurize() {
                                   << " from block " << block << "\n");
     if (!isFnEntryBlock(block)) {
       for (BlockArgument blockArg : block->getArguments()) {
-        auto newArg = newBlock->addArgument(blockArg.getType());
+        auto newArg =
+            newBlock->addArgument(blockArg.getType(), blockArg.getLoc());
         mapper.map(blockArg, newArg);
         LLVM_DEBUG(logger.startLine() << "[cf] remapped block argument "
                                       << blockArg << " to " << newArg << "\n");
@@ -1797,9 +1798,8 @@ LogicalResult ControlFlowStructurizer::structurize() {
     // The loop header block may have block arguments. Since now we place the
     // loop op inside the old merge block, we need to make sure the old merge
     // block has the same block argument list.
-    for (BlockArgument blockArg : headerBlock->getArguments()) {
-      mergeBlock->addArgument(blockArg.getType());
-    }
+    for (BlockArgument blockArg : headerBlock->getArguments())
+      mergeBlock->addArgument(blockArg.getType(), blockArg.getLoc());
 
     // If the loop header block has block arguments, make sure the spv.Branch op
     // matches.

diff  --git a/mlir/lib/Transforms/BufferResultsToOutParams.cpp b/mlir/lib/Transforms/BufferResultsToOutParams.cpp
index 2d74c4085e70f..c462955d068a4 100644
--- a/mlir/lib/Transforms/BufferResultsToOutParams.cpp
+++ b/mlir/lib/Transforms/BufferResultsToOutParams.cpp
@@ -50,8 +50,9 @@ static void updateFuncOp(FuncOp func,
   // Add the new arguments to the entry block if the function is not external.
   if (func.isExternal())
     return;
-  auto newArgs = func.front().addArguments(erasedResultTypes);
-  appendedEntryArgs.append(newArgs.begin(), newArgs.end());
+  Location loc = func.getLoc();
+  for (Type type : erasedResultTypes)
+    appendedEntryArgs.push_back(func.front().addArgument(type, loc));
 }
 
 // Updates all ReturnOps in the scope of the given FuncOp by either keeping them

diff  --git a/mlir/lib/Transforms/NormalizeMemRefs.cpp b/mlir/lib/Transforms/NormalizeMemRefs.cpp
index a27a3d73fcd92..5119c4526364a 100644
--- a/mlir/lib/Transforms/NormalizeMemRefs.cpp
+++ b/mlir/lib/Transforms/NormalizeMemRefs.cpp
@@ -332,6 +332,8 @@ void NormalizeMemRefs::normalizeFuncOpMemRefs(FuncOp funcOp,
   OpBuilder b(funcOp);
 
   FunctionType functionType = funcOp.getType();
+  SmallVector<Location> functionArgLocs(llvm::map_range(
+      funcOp.getArguments(), [](BlockArgument arg) { return arg.getLoc(); }));
   SmallVector<Type, 8> inputTypes;
   // Walk over each argument of a function to perform memref normalization (if
   for (unsigned argIndex :
@@ -356,8 +358,8 @@ void NormalizeMemRefs::normalizeFuncOpMemRefs(FuncOp funcOp,
     }
 
     // Insert a new temporary argument with the new memref type.
-    BlockArgument newMemRef =
-        funcOp.front().insertArgument(argIndex, newMemRefType);
+    BlockArgument newMemRef = funcOp.front().insertArgument(
+        argIndex, newMemRefType, functionArgLocs[argIndex]);
     BlockArgument oldMemRef = funcOp.getArgument(argIndex + 1);
     AffineMap layoutMap = memrefType.getLayout().getAffineMap();
     // Replace all uses of the old memref.

diff  --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index 0bfcf560c353a..cbc8532606281 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -763,7 +763,11 @@ Block *ArgConverter::applySignatureConversion(
   Block *newBlock = block->splitBlock(block->begin());
   block->replaceAllUsesWith(newBlock);
 
-  SmallVector<Value, 4> newArgRange(newBlock->addArguments(convertedTypes));
+  // FIXME: We should map the new arguments to proper locations.
+  SmallVector<Location> newLocs(convertedTypes.size(),
+                                rewriter.getUnknownLoc());
+  SmallVector<Value, 4> newArgRange(
+      newBlock->addArguments(convertedTypes, newLocs));
   ArrayRef<Value> newArgs(newArgRange);
 
   // Remap each of the original arguments as determined by the signature

diff  --git a/mlir/lib/Transforms/Utils/InliningUtils.cpp b/mlir/lib/Transforms/Utils/InliningUtils.cpp
index 8b2040633a1ae..2a4c5dab95341 100644
--- a/mlir/lib/Transforms/Utils/InliningUtils.cpp
+++ b/mlir/lib/Transforms/Utils/InliningUtils.cpp
@@ -216,8 +216,9 @@ inlineRegionImpl(InlinerInterface &interface, Region *src, Block *inlineBlock,
     // Otherwise, there were multiple blocks inlined. Add arguments to the post
     // insertion block to represent the results to replace.
     for (const auto &resultToRepl : llvm::enumerate(resultsToReplace)) {
-      resultToRepl.value().replaceAllUsesWith(postInsertBlock->addArgument(
-          regionResultTypes[resultToRepl.index()]));
+      resultToRepl.value().replaceAllUsesWith(
+          postInsertBlock->addArgument(regionResultTypes[resultToRepl.index()],
+                                       resultToRepl.value().getLoc()));
     }
 
     /// Handle the terminators for each of the new blocks.

diff  --git a/mlir/lib/Transforms/Utils/RegionUtils.cpp b/mlir/lib/Transforms/Utils/RegionUtils.cpp
index e60e7e65cd603..cf7761073d07a 100644
--- a/mlir/lib/Transforms/Utils/RegionUtils.cpp
+++ b/mlir/lib/Transforms/Utils/RegionUtils.cpp
@@ -601,8 +601,11 @@ LogicalResult BlockMergeCluster::merge(RewriterBase &rewriter) {
         newArguments[i][it.index()] = operand.get();
 
         // Update the operand and insert an argument if this is the leader.
-        if (i == 0)
-          operand.set(leaderBlock->addArgument(operand.get().getType()));
+        if (i == 0) {
+          Value operandVal = operand.get();
+          operand.set(leaderBlock->addArgument(operandVal.getType(),
+                                               operandVal.getLoc()));
+        }
       }
     }
     // Update the predecessors for each of the blocks.

diff  --git a/mlir/test/CAPI/ir.c b/mlir/test/CAPI/ir.c
index 418d5a0aa3fe8..d01ccaeb0e93a 100644
--- a/mlir/test/CAPI/ir.c
+++ b/mlir/test/CAPI/ir.c
@@ -81,9 +81,11 @@ MlirModule makeAndDumpAdd(MlirContext ctx, MlirLocation location) {
   MlirType memrefType =
       mlirTypeParseGet(ctx, mlirStringRefCreateFromCString("memref<?xf32>"));
   MlirType funcBodyArgTypes[] = {memrefType, memrefType};
+  MlirLocation funcBodyArgLocs[] = {location, location};
   MlirRegion funcBodyRegion = mlirRegionCreate();
-  MlirBlock funcBody = mlirBlockCreate(
-      sizeof(funcBodyArgTypes) / sizeof(MlirType), funcBodyArgTypes);
+  MlirBlock funcBody =
+      mlirBlockCreate(sizeof(funcBodyArgTypes) / sizeof(MlirType),
+                      funcBodyArgTypes, funcBodyArgLocs);
   mlirRegionAppendOwnedBlock(funcBodyRegion, funcBody);
 
   MlirAttribute funcTypeAttr = mlirAttributeParseGet(
@@ -130,8 +132,8 @@ MlirModule makeAndDumpAdd(MlirContext ctx, MlirLocation location) {
   mlirBlockAppendOwnedOperation(funcBody, dim);
 
   MlirRegion loopBodyRegion = mlirRegionCreate();
-  MlirBlock loopBody = mlirBlockCreate(0, NULL);
-  mlirBlockAddArgument(loopBody, indexType);
+  MlirBlock loopBody = mlirBlockCreate(0, NULL, NULL);
+  mlirBlockAddArgument(loopBody, indexType, location);
   mlirRegionAppendOwnedBlock(loopBodyRegion, loopBody);
 
   MlirAttribute indexOneLiteral =
@@ -507,10 +509,10 @@ static void buildWithInsertionsAndPrint(MlirContext ctx) {
   MlirType i2 = mlirIntegerTypeGet(ctx, 2);
   MlirType i3 = mlirIntegerTypeGet(ctx, 3);
   MlirType i4 = mlirIntegerTypeGet(ctx, 4);
-  MlirBlock block1 = mlirBlockCreate(1, &i1);
-  MlirBlock block2 = mlirBlockCreate(1, &i2);
-  MlirBlock block3 = mlirBlockCreate(1, &i3);
-  MlirBlock block4 = mlirBlockCreate(1, &i4);
+  MlirBlock block1 = mlirBlockCreate(1, &i1, &loc);
+  MlirBlock block2 = mlirBlockCreate(1, &i2, &loc);
+  MlirBlock block3 = mlirBlockCreate(1, &i3, &loc);
+  MlirBlock block4 = mlirBlockCreate(1, &i4, &loc);
   // Insert blocks so as to obtain the 1-2-3-4 order,
   mlirRegionInsertOwnedBlockBefore(region, nullBlock, block3);
   mlirRegionInsertOwnedBlockBefore(region, block3, block2);
@@ -1549,7 +1551,7 @@ static int testBackreferences() {
   MlirOperationState opState =
       mlirOperationStateGet(mlirStringRefCreateFromCString("invalid.op"), loc);
   MlirRegion region = mlirRegionCreate();
-  MlirBlock block = mlirBlockCreate(0, NULL);
+  MlirBlock block = mlirBlockCreate(0, NULL, NULL);
   mlirRegionAppendOwnedBlock(region, block);
   mlirOperationStateAddOwnedRegions(&opState, 1, &region);
   MlirOperation op = mlirOperationCreate(&opState);

diff  --git a/mlir/test/lib/Dialect/Test/TestPatterns.cpp b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
index 7e20b13b8ef81..34dd14176b45d 100644
--- a/mlir/test/lib/Dialect/Test/TestPatterns.cpp
+++ b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
@@ -297,7 +297,8 @@ struct TestRegionRewriteUndo : public RewritePattern {
     newRegion.addRegion();
     auto *regionOp = rewriter.createOperation(newRegion);
     auto *entryBlock = rewriter.createBlock(&regionOp->getRegion(0));
-    entryBlock->addArgument(rewriter.getIntegerType(64));
+    entryBlock->addArgument(rewriter.getIntegerType(64),
+                            rewriter.getUnknownLoc());
 
     // Add an explicitly illegal operation to ensure the conversion fails.
     rewriter.create<ILLegalOpF>(op->getLoc(), rewriter.getIntegerType(32));
@@ -318,8 +319,9 @@ struct TestCreateBlock : public RewritePattern {
                                 PatternRewriter &rewriter) const final {
     Region &region = *op->getParentRegion();
     Type i32Type = rewriter.getIntegerType(32);
-    rewriter.createBlock(&region, region.end(), {i32Type, i32Type});
-    rewriter.create<TerminatorOp>(op->getLoc());
+    Location loc = op->getLoc();
+    rewriter.createBlock(&region, region.end(), {i32Type, i32Type}, {loc, loc});
+    rewriter.create<TerminatorOp>(loc);
     rewriter.replaceOp(op, {});
     return success();
   }
@@ -335,10 +337,11 @@ struct TestCreateIllegalBlock : public RewritePattern {
                                 PatternRewriter &rewriter) const final {
     Region &region = *op->getParentRegion();
     Type i32Type = rewriter.getIntegerType(32);
-    rewriter.createBlock(&region, region.end(), {i32Type, i32Type});
+    Location loc = op->getLoc();
+    rewriter.createBlock(&region, region.end(), {i32Type, i32Type}, {loc, loc});
     // Create an illegal op to ensure the conversion fails.
-    rewriter.create<ILLegalOpF>(op->getLoc(), i32Type);
-    rewriter.create<TerminatorOp>(op->getLoc());
+    rewriter.create<ILLegalOpF>(loc, i32Type);
+    rewriter.create<TerminatorOp>(loc);
     rewriter.replaceOp(op, {});
     return success();
   }

diff  --git a/mlir/test/lib/IR/TestFunc.cpp b/mlir/test/lib/IR/TestFunc.cpp
index 7ad20756f6f16..dee9f8a5b2e52 100644
--- a/mlir/test/lib/IR/TestFunc.cpp
+++ b/mlir/test/lib/IR/TestFunc.cpp
@@ -20,6 +20,7 @@ struct TestFuncInsertArg
   void runOnOperation() override {
     auto module = getOperation();
 
+    UnknownLoc unknownLoc = UnknownLoc::get(module.getContext());
     for (FuncOp func : module.getOps<FuncOp>()) {
       auto inserts = func->getAttrOfType<ArrayAttr>("test.insert_args");
       if (!inserts || inserts.empty())
@@ -27,7 +28,7 @@ struct TestFuncInsertArg
       SmallVector<unsigned, 4> indicesToInsert;
       SmallVector<Type, 4> typesToInsert;
       SmallVector<DictionaryAttr, 4> attrsToInsert;
-      SmallVector<Optional<Location>, 4> locsToInsert;
+      SmallVector<Location, 4> locsToInsert;
       for (auto insert : inserts.getAsRange<ArrayAttr>()) {
         indicesToInsert.push_back(
             insert[0].cast<IntegerAttr>().getValue().getZExtValue());
@@ -35,10 +36,9 @@ struct TestFuncInsertArg
         attrsToInsert.push_back(insert.size() > 2
                                     ? insert[2].cast<DictionaryAttr>()
                                     : DictionaryAttr::get(&getContext()));
-        locsToInsert.push_back(
-            insert.size() > 3
-                ? Optional<Location>(insert[3].cast<LocationAttr>())
-                : Optional<Location>{});
+        locsToInsert.push_back(insert.size() > 3
+                                   ? Location(insert[3].cast<LocationAttr>())
+                                   : unknownLoc);
       }
       func->removeAttr("test.insert_args");
       func.insertArguments(indicesToInsert, typesToInsert, attrsToInsert,


        


More information about the flang-commits mailing list