[IR] Add an ImplicitLocOpBuilder helper class for building IR with the same loc.

[lattner]

Summary

One common situation is to create a lot of IR at a well known location,
e.g. when doing a big rewrite from one dialect to another where you're expanding
ops out into lots of other ops.

For these sorts of situations, it is annoying to pass the location into
every create call. As we discused in a few threads on the forum, a way to help
with this is to produce a new sort of builder that holds a location and provides
it to each of the create<> calls automatically.

This patch implements an ImplicitLocOpBuilder class that does this. We've had
good experience with this in the CIRCT project, and it makes sense to upstream to
MLIR.

I picked a random pass to adopt it to show the impact, but I don't think there is
any particular need to force adopt it in the codebase.

[ftynse]

Looks great!

mlir/include/mlir/IR/ImplicitLocOpBuilder.h

2

Please add the license header

19–22

Are these necessary? We are now in the mlir:: namespace.

34

std::forward<T>(operand)?

[lattner]

Fixes AlexZ recommended.

[lattner]

Great catches, fixed. Thanks!

[bondhugula]

mlir/include/mlir/IR/ImplicitLocOpBuilder.h

22

It would have been good to mention that this should be passed by reference? With the location, this would have four fields.

118

This is missing a doc comment.

[lattner]

mlir/include/mlir/IR/ImplicitLocOpBuilder.h

22

I think this should be passed the same way as any builder. Builders that are copied take their state with them. If you pass by reference you maintain it. It is reasonable to splice off the OpBuilder subset of this (by value) if that is the desired behavior.

mlir/include/mlir/IR/ImplicitLocOpBuilder.h

+//===- ImplicitLocOpBuilder.h - Convenience OpBuilder -----------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Helper class to create ops with a modally set location.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_IR_IMPLICITLOCOPBUILDER_H
+#define MLIR_IR_IMPLICITLOCOPBUILDER_H
+
+#include "mlir/IR/Builders.h"
+
+namespace mlir {
+
+/// ImplictLocOpBuilder maintains a 'current location', allowing use of the
+/// create<> method without specifying the location.  It is otherwise the same
+/// as OpBuilder.
+class ImplicitLocOpBuilder : public mlir::OpBuilder {
+public:
+  /// Create an ImplicitLocOpBuilder using the insertion point and listener from
+  /// an existing OpBuilder.
+  ImplicitLocOpBuilder(Location loc, const OpBuilder &builder)
+      : OpBuilder(builder), curLoc(loc) {}
+
+  /// OpBuilder has a bunch of convenience constructors - we support them all
+  /// with the additional Location.
+  template <typename T>
+  ImplicitLocOpBuilder(Location loc, T &&operand, Listener *listener = nullptr)
+      : OpBuilder(std::forward<T>(operand), listener), curLoc(loc) {}
+
+  ImplicitLocOpBuilder(Location loc, Block *block, Block::iterator insertPoint,
+                       Listener *listener = nullptr)
+      : OpBuilder(block, insertPoint, listener), curLoc(loc) {}
+
+  /// Create a builder and set the insertion point to before the first operation
+  /// in the block but still inside the block.
+  static ImplicitLocOpBuilder atBlockBegin(Location loc, Block *block,
+                                           Listener *listener = nullptr) {
+    return ImplicitLocOpBuilder(loc, block, block->begin(), listener);
+  }
+
+  /// Create a builder and set the insertion point to after the last operation
+  /// in the block but still inside the block.
+  static ImplicitLocOpBuilder atBlockEnd(Location loc, Block *block,
+                                         Listener *listener = nullptr) {
+    return ImplicitLocOpBuilder(loc, block, block->end(), listener);
+  }
+
+  /// Create a builder and set the insertion point to before the block
+  /// terminator.
+  static ImplicitLocOpBuilder atBlockTerminator(Location loc, Block *block,
+                                                Listener *listener = nullptr) {
+    auto *terminator = block->getTerminator();
+    assert(terminator != nullptr && "the block has no terminator");
+    return ImplicitLocOpBuilder(loc, block, Block::iterator(terminator),
+                                listener);
+  }
+
+  /// Accessors for the implied location.
+  Location getLoc() const { return curLoc; }
+  void setLoc(Location loc) { curLoc = loc; }
+
+  // We allow clients to use the explicit-loc version of create as well.
+  using OpBuilder::create;
+  using OpBuilder::createOrFold;
+
+  /// Create an operation of specific op type at the current insertion point and
+  /// location.
+  template <typename OpTy, typename... Args>
+  OpTy create(Args &&... args) {
+    return OpBuilder::create<OpTy>(curLoc, std::forward<Args>(args)...);
+  }
+
+  /// Create an operation of specific op type at the current insertion point,
+  /// and immediately try to fold it. This functions populates 'results' with
+  /// the results after folding the operation.
+  template <typename OpTy, typename... Args>
+  void createOrFold(llvm::SmallVectorImpl<Value> &results, Args &&... args) {
+    OpBuilder::createOrFold<OpTy>(results, curLoc, std::forward<Args>(args)...);
+  }
+
+  /// Overload to create or fold a single result operation.
+  template <typename OpTy, typename... Args>
+  typename std::enable_if<OpTy::template hasTrait<mlir::OpTrait::OneResult>(),
+                          Value>::type
+  createOrFold(Args &&... args) {
+    return OpBuilder::createOrFold<OpTy>(curLoc, std::forward<Args>(args)...);
+  }
+
+  /// Overload to create or fold a zero result operation.
+  template <typename OpTy, typename... Args>
+  typename std::enable_if<OpTy::template hasTrait<mlir::OpTrait::ZeroResult>(),
+                          OpTy>::type
+  createOrFold(Args &&... args) {
+    return OpBuilder::createOrFold<OpTy>(curLoc, std::forward<Args>(args)...);
+  }
+
+  /// This builder can also be used to emit diagnostics to the current location.
+  mlir::InFlightDiagnostic
+  emitError(const llvm::Twine &message = llvm::Twine()) {
+    return mlir::emitError(curLoc, message);
+  }
+  mlir::InFlightDiagnostic
+  emitWarning(const llvm::Twine &message = llvm::Twine()) {
+    return mlir::emitWarning(curLoc, message);
+  }
+  mlir::InFlightDiagnostic
+  emitRemark(const llvm::Twine &message = llvm::Twine()) {
+    return mlir::emitRemark(curLoc, message);
+  }
+
+private:
+  Location curLoc;
+};
+
+} // namespace mlir
+
+#endif // MLIR_IR_IMPLICITLOCOPBUILDER_H
 No newline at end of file

mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp

 //===- AsyncToLLVM.cpp - Convert Async to LLVM dialect --------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/AsyncToLLVM/AsyncToLLVM.h"
 
 #include "../PassDetail.h"
 #include "mlir/Dialect/Async/IR/Async.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/BlockAndValueMapping.h"
-#include "mlir/IR/Builders.h"
+#include "mlir/IR/ImplicitLocOpBuilder.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/RegionUtils.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/Support/FormatVariadic.h"
 
 #define DEBUG_TYPE "convert-async-to-llvm"
static LLVM::LLVMType resumeFunctionType(MLIRContext *ctx) {
     auto i8Ptr = LLVM::LLVMType::getInt8PtrTy(ctx);
     return LLVM::LLVMType::getFunctionTy(voidTy, {i8Ptr}, false);
   }
 };
 } // namespace
 
 // Adds Async Runtime C API declarations to the module.
 static void addAsyncRuntimeApiDeclarations(ModuleOp module) {
-  auto builder = OpBuilder::atBlockTerminator(module.getBody());
+  auto builder = ImplicitLocOpBuilder::atBlockTerminator(module.getLoc(),
+                                                         module.getBody());
 
   auto addFuncDecl = [&](StringRef name, FunctionType type) {
     if (module.lookupSymbol(name))
       return;
-    builder.create<FuncOp>(module.getLoc(), name, type).setPrivate();
+    builder.create<FuncOp>(name, type).setPrivate();
   };
 
   MLIRContext *ctx = module.getContext();
   addFuncDecl(kAddRef, AsyncAPI::addOrDropRefFunctionType(ctx));
   addFuncDecl(kDropRef, AsyncAPI::addOrDropRefFunctionType(ctx));
   addFuncDecl(kCreateToken, AsyncAPI::createTokenFunctionType(ctx));
   addFuncDecl(kCreateGroup, AsyncAPI::createGroupFunctionType(ctx));
   addFuncDecl(kEmplaceToken, AsyncAPI::emplaceTokenFunctionType(ctx));

 static constexpr const char *kCoroBegin = "llvm.coro.begin";
 static constexpr const char *kCoroSave = "llvm.coro.save";
 static constexpr const char *kCoroSuspend = "llvm.coro.suspend";
 static constexpr const char *kCoroEnd = "llvm.coro.end";
 static constexpr const char *kCoroFree = "llvm.coro.free";
 static constexpr const char *kCoroResume = "llvm.coro.resume";
 
 /// Adds an LLVM function declaration to a module.
-static void addLLVMFuncDecl(ModuleOp module, OpBuilder &builder, StringRef name,
-                            LLVM::LLVMType ret,
+static void addLLVMFuncDecl(ModuleOp module, ImplicitLocOpBuilder &builder,
+                            StringRef name, LLVM::LLVMType ret,
                             ArrayRef<LLVM::LLVMType> params) {
   if (module.lookupSymbol(name))
     return;
   LLVM::LLVMType type = LLVM::LLVMType::getFunctionTy(ret, params, false);
-  builder.create<LLVM::LLVMFuncOp>(module.getLoc(), name, type);
+  builder.create<LLVM::LLVMFuncOp>(name, type);
 }
 
 /// Adds coroutine intrinsics declarations to the module.
 static void addCoroutineIntrinsicsDeclarations(ModuleOp module) {
   using namespace mlir::LLVM;
 
   MLIRContext *ctx = module.getContext();
-  OpBuilder builder(module.getBody()->getTerminator());
+  ImplicitLocOpBuilder builder(module.getLoc(),
+                               module.getBody()->getTerminator());
 
   auto token = LLVMTokenType::get(ctx);
   auto voidTy = LLVMType::getVoidTy(ctx);
 
   auto i8 = LLVMType::getInt8Ty(ctx);
   auto i1 = LLVMType::getInt1Ty(ctx);
   auto i32 = LLVMType::getInt32Ty(ctx);
   auto i64 = LLVMType::getInt64Ty(ctx);

 static constexpr const char *kMalloc = "malloc";
 static constexpr const char *kFree = "free";
 
 /// Adds malloc/free declarations to the module.
 static void addCRuntimeDeclarations(ModuleOp module) {
   using namespace mlir::LLVM;
 
   MLIRContext *ctx = module.getContext();
-  OpBuilder builder(module.getBody()->getTerminator());
+  ImplicitLocOpBuilder builder(module.getLoc(),
+                               module.getBody()->getTerminator());
 
   auto voidTy = LLVMType::getVoidTy(ctx);
   auto i64 = LLVMType::getInt64Ty(ctx);
   auto i8Ptr = LLVMType::getInt8PtrTy(ctx);
 
   addLLVMFuncDecl(module, builder, kMalloc, i8Ptr, {i64});
   addLLVMFuncDecl(module, builder, kFree, voidTy, {i8Ptr});
 }
static void addResumeFunction(ModuleOp module) {
   auto voidTy = LLVM::LLVMType::getVoidTy(ctx);
   auto i8Ptr = LLVM::LLVMType::getInt8PtrTy(ctx);
 
   auto resumeOp = moduleBuilder.create<LLVM::LLVMFuncOp>(
       loc, kResume, LLVM::LLVMType::getFunctionTy(voidTy, {i8Ptr}, false));
   resumeOp.setPrivate();
 
   auto *block = resumeOp.addEntryBlock();
-  OpBuilder blockBuilder = OpBuilder::atBlockEnd(block);
+  auto blockBuilder = ImplicitLocOpBuilder::atBlockEnd(loc, block);
 
-  blockBuilder.create<LLVM::CallOp>(loc, TypeRange(),
+  blockBuilder.create<LLVM::CallOp>(TypeRange(),
                                     blockBuilder.getSymbolRefAttr(kCoroResume),
                                     resumeOp.getArgument(0));
 
-  blockBuilder.create<LLVM::ReturnOp>(loc, ValueRange());
+  blockBuilder.create<LLVM::ReturnOp>(ValueRange());
 }
 
 //===----------------------------------------------------------------------===//
 // async.execute op outlining to the coroutine functions.
 //===----------------------------------------------------------------------===//
 
 // Function targeted for coroutine transformation has two additional blocks at
 // the end: coroutine cleanup and coroutine suspension.
static CoroMachinery setupCoroMachinery(FuncOp func) {
   auto i1 = LLVM::LLVMType::getInt1Ty(ctx);
   auto i32 = LLVM::LLVMType::getInt32Ty(ctx);
   auto i64 = LLVM::LLVMType::getInt64Ty(ctx);
   auto i8Ptr = LLVM::LLVMType::getInt8PtrTy(ctx);
 
   Block *entryBlock = func.addEntryBlock();
   Location loc = func.getBody().getLoc();
 
-  OpBuilder builder = OpBuilder::atBlockBegin(entryBlock);
+  auto builder = ImplicitLocOpBuilder::atBlockBegin(loc, entryBlock);
 
   // ------------------------------------------------------------------------ //
   // Allocate async tokens/values that we will return from a ramp function.
   // ------------------------------------------------------------------------ //
-  auto createToken =
-      builder.create<CallOp>(loc, kCreateToken, TokenType::get(ctx));
+  auto createToken = builder.create<CallOp>(kCreateToken, TokenType::get(ctx));
 
   // ------------------------------------------------------------------------ //
   // Initialize coroutine: allocate frame, get coroutine handle.
   // ------------------------------------------------------------------------ //
 
   // Constants for initializing coroutine frame.
   auto constZero =
-      builder.create<LLVM::ConstantOp>(loc, i32, builder.getI32IntegerAttr(0));
+      builder.create<LLVM::ConstantOp>(i32, builder.getI32IntegerAttr(0));
   auto constFalse =
-      builder.create<LLVM::ConstantOp>(loc, i1, builder.getBoolAttr(false));
-  auto nullPtr = builder.create<LLVM::NullOp>(loc, i8Ptr);
+      builder.create<LLVM::ConstantOp>(i1, builder.getBoolAttr(false));
+  auto nullPtr = builder.create<LLVM::NullOp>(i8Ptr);
 
   // Get coroutine id: @llvm.coro.id
   auto coroId = builder.create<LLVM::CallOp>(
-      loc, token, builder.getSymbolRefAttr(kCoroId),
+      token, builder.getSymbolRefAttr(kCoroId),
       ValueRange({constZero, nullPtr, nullPtr, nullPtr}));
 
   // Get coroutine frame size: @llvm.coro.size.i64
   auto coroSize = builder.create<LLVM::CallOp>(
-      loc, i64, builder.getSymbolRefAttr(kCoroSizeI64), ValueRange());
+      i64, builder.getSymbolRefAttr(kCoroSizeI64), ValueRange());
 
   // Allocate memory for coroutine frame.
-  auto coroAlloc = builder.create<LLVM::CallOp>(
-      loc, i8Ptr, builder.getSymbolRefAttr(kMalloc),
+  auto coroAlloc =
+      builder.create<LLVM::CallOp>(i8Ptr, builder.getSymbolRefAttr(kMalloc),
                                    ValueRange(coroSize.getResult(0)));
 
   // Begin a coroutine: @llvm.coro.begin
   auto coroHdl = builder.create<LLVM::CallOp>(
-      loc, i8Ptr, builder.getSymbolRefAttr(kCoroBegin),
+      i8Ptr, builder.getSymbolRefAttr(kCoroBegin),
       ValueRange({coroId.getResult(0), coroAlloc.getResult(0)}));
 
   Block *cleanupBlock = func.addBlock();
   Block *suspendBlock = func.addBlock();
 
   // ------------------------------------------------------------------------ //
   // Coroutine cleanup block: deallocate coroutine frame, free the memory.
   // ------------------------------------------------------------------------ //
   builder.setInsertionPointToStart(cleanupBlock);
 
   // Get a pointer to the coroutine frame memory: @llvm.coro.free.
   auto coroMem = builder.create<LLVM::CallOp>(
-      loc, i8Ptr, builder.getSymbolRefAttr(kCoroFree),
+      i8Ptr, builder.getSymbolRefAttr(kCoroFree),
       ValueRange({coroId.getResult(0), coroHdl.getResult(0)}));
 
   // Free the memory.
-  builder.create<LLVM::CallOp>(loc, TypeRange(),
-                               builder.getSymbolRefAttr(kFree),
+  builder.create<LLVM::CallOp>(TypeRange(), builder.getSymbolRefAttr(kFree),
                                ValueRange(coroMem.getResult(0)));
   // Branch into the suspend block.
-  builder.create<BranchOp>(loc, suspendBlock);
+  builder.create<BranchOp>(suspendBlock);
 
   // ------------------------------------------------------------------------ //
   // Coroutine suspend block: mark the end of a coroutine and return allocated
   // async token.
   // ------------------------------------------------------------------------ //
   builder.setInsertionPointToStart(suspendBlock);
 
   // Mark the end of a coroutine: @llvm.coro.end.
-  builder.create<LLVM::CallOp>(loc, i1, builder.getSymbolRefAttr(kCoroEnd),
+  builder.create<LLVM::CallOp>(i1, builder.getSymbolRefAttr(kCoroEnd),
                                ValueRange({coroHdl.getResult(0), constFalse}));
 
   // Return created `async.token` from the suspend block. This will be the
   // return value of a coroutine ramp function.
-  builder.create<ReturnOp>(loc, createToken.getResult(0));
+  builder.create<ReturnOp>(createToken.getResult(0));
 
   // Branch from the entry block to the cleanup block to create a valid CFG.
   builder.setInsertionPointToEnd(entryBlock);
 
-  builder.create<BranchOp>(loc, cleanupBlock);
+  builder.create<BranchOp>(cleanupBlock);
 
   // `async.await` op lowering will create resume blocks for async
   // continuations, and will conditionally branch to cleanup or suspend blocks.
 
   return {createToken.getResult(0), coroHdl.getResult(0), cleanupBlock,
           suspendBlock};
 }
 

 // Note that this is not reversible transformation.
 static std::pair<FuncOp, CoroMachinery>
 outlineExecuteOp(SymbolTable &symbolTable, ExecuteOp execute) {
   ModuleOp module = execute->getParentOfType<ModuleOp>();
 
   MLIRContext *ctx = module.getContext();
   Location loc = execute.getLoc();
 
-  OpBuilder moduleBuilder(module.getBody()->getTerminator());
-
   // Collect all outlined function inputs.
   llvm::SetVector<mlir::Value> functionInputs(execute.dependencies().begin(),
                                               execute.dependencies().end());
   getUsedValuesDefinedAbove(execute.body(), functionInputs);
 
   // Collect types for the outlined function inputs and outputs.
   auto typesRange = llvm::map_range(
       functionInputs, [](Value value) { return value.getType(); });
   SmallVector<Type, 4> inputTypes(typesRange.begin(), typesRange.end());
   auto outputTypes = execute.getResultTypes();
 
-  auto funcType = moduleBuilder.getFunctionType(inputTypes, outputTypes);
+  auto funcType = FunctionType::get(ctx, inputTypes, outputTypes);
   auto funcAttrs = ArrayRef<NamedAttribute>();
 
   // TODO: Derive outlined function name from the parent FuncOp (support
   // multiple nested async.execute operations).
   FuncOp func = FuncOp::create(loc, kAsyncFnPrefix, funcType, funcAttrs);
-  symbolTable.insert(func, moduleBuilder.getInsertionPoint());
+  symbolTable.insert(func, Block::iterator(module.getBody()->getTerminator()));
 
   SymbolTable::setSymbolVisibility(func, SymbolTable::Visibility::Private);
 
   // Prepare a function for coroutine lowering by adding entry/cleanup/suspend
   // blocks, adding llvm.coro instrinsics and setting up control flow.
   CoroMachinery coro = setupCoroMachinery(func);
 
   // Suspend async function at the end of an entry block, and resume it using
   // Async execute API (execution will be resumed in a thread managed by the
   // async runtime).
   Block *entryBlock = &func.getBlocks().front();
-  OpBuilder builder = OpBuilder::atBlockTerminator(entryBlock);
+  auto builder = ImplicitLocOpBuilder::atBlockTerminator(loc, entryBlock);
 
   // A pointer to coroutine resume intrinsic wrapper.
   auto resumeFnTy = AsyncAPI::resumeFunctionType(ctx);
   auto resumePtr = builder.create<LLVM::AddressOfOp>(
-      loc, LLVM::LLVMPointerType::get(resumeFnTy), kResume);
+      LLVM::LLVMPointerType::get(resumeFnTy), kResume);
 
   // Save the coroutine state: @llvm.coro.save
   auto coroSave = builder.create<LLVM::CallOp>(
-      loc, LLVM::LLVMTokenType::get(ctx), builder.getSymbolRefAttr(kCoroSave),
+      LLVM::LLVMTokenType::get(ctx), builder.getSymbolRefAttr(kCoroSave),
       ValueRange({coro.coroHandle}));
 
   // Call async runtime API to execute a coroutine in the managed thread.
   SmallVector<Value, 2> executeArgs = {coro.coroHandle, resumePtr.res()};
-  builder.create<CallOp>(loc, TypeRange(), kExecute, executeArgs);
+  builder.create<CallOp>(TypeRange(), kExecute, executeArgs);
 
   // Split the entry block before the terminator.
   auto *terminatorOp = entryBlock->getTerminator();
   Block *suspended = terminatorOp->getBlock();
   Block *resume = suspended->splitBlock(terminatorOp);
   addSuspensionPoint(coro, coroSave.getResult(0), terminatorOp, suspended,
                      resume, builder);
 
   // Await on all dependencies before starting to execute the body region.
   builder.setInsertionPointToStart(resume);
   for (size_t i = 0; i < execute.dependencies().size(); ++i)
-    builder.create<AwaitOp>(loc, func.getArgument(i));
+    builder.create<AwaitOp>(func.getArgument(i));
 
   // Map from function inputs defined above the execute op to the function
   // arguments.
   BlockAndValueMapping valueMapping;
   valueMapping.map(functionInputs, func.getArguments());
 
   // Clone all operations from the execute operation body into the outlined
   // function body, and replace all `async.yield` operations with a call
   // to async runtime to emplace the result token.
   for (Operation &op : execute.body().getOps()) {
     if (isa<async::YieldOp>(op)) {
-      builder.create<CallOp>(loc, kEmplaceToken, TypeRange(), coro.asyncToken);
+      builder.create<CallOp>(kEmplaceToken, TypeRange(), coro.asyncToken);
       continue;
     }
     builder.clone(op, valueMapping);
   }
 
   // Replace the original `async.execute` with a call to outlined function.
-  OpBuilder callBuilder(execute);
-  auto callOutlinedFunc =
-      callBuilder.create<CallOp>(loc, func.getName(), execute.getResultTypes(),
-                                 functionInputs.getArrayRef());
+  ImplicitLocOpBuilder callBuilder(loc, execute);
+  auto callOutlinedFunc = callBuilder.create<CallOp>(
+      func.getName(), execute.getResultTypes(), functionInputs.getArrayRef());
   execute.replaceAllUsesWith(callOutlinedFunc.getResults());
   execute.erase();
 
   return {func, coro};
 }
 
 //===----------------------------------------------------------------------===//
 // Convert Async dialect types to LLVM types.
if (!isInCoroutine)
       rewriter.create<CallOp>(loc, TypeRange(), blockingAwaitFuncName,
                               ValueRange(operands[0]));
 
     // Inside the coroutine we convert await operation into coroutine suspension
     // point, and resume execution asynchronously.
     if (isInCoroutine) {
       const CoroMachinery &coro = outlined->getSecond();
 
-      OpBuilder builder(op, rewriter.getListener());
+      ImplicitLocOpBuilder builder(loc, op, rewriter.getListener());
       MLIRContext *ctx = op->getContext();
 
       // A pointer to coroutine resume intrinsic wrapper.
       auto resumeFnTy = AsyncAPI::resumeFunctionType(ctx);
       auto resumePtr = builder.create<LLVM::AddressOfOp>(
-          loc, LLVM::LLVMPointerType::get(resumeFnTy), kResume);
+          LLVM::LLVMPointerType::get(resumeFnTy), kResume);
 
       // Save the coroutine state: @llvm.coro.save
       auto coroSave = builder.create<LLVM::CallOp>(
-          loc, LLVM::LLVMTokenType::get(ctx),
-          builder.getSymbolRefAttr(kCoroSave), ValueRange(coro.coroHandle));
+          LLVM::LLVMTokenType::get(ctx), builder.getSymbolRefAttr(kCoroSave),
+          ValueRange(coro.coroHandle));
 
       // Call async runtime API to resume a coroutine in the managed thread when
       // the async await argument becomes ready.
       SmallVector<Value, 3> awaitAndExecuteArgs = {operands[0], coro.coroHandle,
                                                    resumePtr.res()};
-      builder.create<CallOp>(loc, TypeRange(), coroAwaitFuncName,
+      builder.create<CallOp>(TypeRange(), coroAwaitFuncName,
                              awaitAndExecuteArgs);
 
       Block *suspended = op->getBlock();
 
       // Split the entry block before the await operation.
       Block *resume = rewriter.splitBlock(suspended, Block::iterator(op));
       addSuspensionPoint(coro, coroSave.getResult(0), op, suspended, resume,
                          builder);