[mlir][gpu] Add the Select Object compilation attribute.

[fmorac]

Summary

For an explanation of these patches see D154153.

Commit message:
This patch adds the default offloading handler for GPU binary ops: #gpu.select_object,
it selects the object to embed based on an index or a target attribute, embedding
the object as a global string and launches the kernel using the scheme used in the
GPU to LLVM pass.

Depends on D154137

[fmorac]

Rebasing.

[krzysz00]

Could this hypothetically be used to do runtime selection - that is, if I'll have a gfx90a binary and a gfx940 binary, could I write an object selector that'll query which GPU I have at execution time and select the relevant binary? Or is that out of scope here?

[fmorac]

In D154147#4465163, @krzysz00 wrote:

Could this hypothetically be used to do runtime selection - that is, if I'll have a gfx90a binary and a gfx940 binary, could I write an object selector that'll query which GPU I have at execution time and select the relevant binary? Or is that out of scope here?

With this particular attribute, no.

However If you add the relevant runtime functions and create your own ObjectManager attribute -see D154108 for the interface, you could definitely do it. It would look something like:

gpu.binary @myobject <#mydialect.runtime_select_object> [object0, object1]

I even think it would be possible to have AMD and NVIDIA targets all packed in a single IR, and perform dispatching based on the GPU at execution time. The ObjectManager attribute interface leaves the room open for all of this.

[mehdi_amini]

Could this hypothetically be used to do runtime selection - that is, if I'll have a gfx90a binary and a gfx940 binary, could I write an object selector that'll query which GPU I have at execution time and select the relevant binary? Or is that out of scope here?

That would make the compiler all but stateless, which is something I would be strongly against.

[fmorac]

In D154147#4465656, @mehdi_amini wrote:

Could this hypothetically be used to do runtime selection - that is, if I'll have a gfx90a binary and a gfx940 binary, could I write an object selector that'll query which GPU I have at execution time and select the relevant binary? Or is that out of scope here?

That would make the compiler all but stateless, which is something I would be strongly against.

I think we were talking about something along the lines of fat binaries. Like implementing a fatbin Object Manager attribute, something like:

gpu.binary @name <#gpu.fatbin> [#gpu.object<#gpu.amdgpu<chip = "gfx90a">, "BLOB">, #gpu.object<#gpu.amdgpu<chip = "gfx940">, "BLOB">, #gpu.object<#gpu.nvptx<chip = "sm_80">, "BLOB">]

That when translated down to LLVM IR embeds all the objects together in a single image. And instead of calling mgpuModuleLoad we have something like mgpuFatbinLoad and mgpuFatbinGetFunction loading the correct function at execution time depending on the detected GPU.

The above scheme is possible with Object Manager attributes.

[krzysz00]

Yeah, I was thinking something more like fat binaries here.

[krzysz00]

Looks reasonable to me

[fmorac]

Rebasing.

[mehdi_amini]

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1112

It's not clear to me when it makes sense to have multiple gpu.object when there is a select?

[fmorac]

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1112

The SelectObjectManager is just the default manager, but you can have FatBinManager in which it makes sense to have multiple objects.

As to why select_object is the way it is, is to provide a sort of similar method to that of gpu-binary-annotation in the existing implementation of gpu-to-llvm where you can choose the binary:

gpu.module @kernel_module attributes {
      nvvm.cubin = "CUBIN", rocdl.hsaco = "HSACO"
  }

[mehdi_amini]

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1112

How would the select be used? The gpu-binary-annotation is an input to a transformation and not an IR construct.

[fmorac]

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

1112

I thought of 2 uses cases:

1. having a pass that lets you specify the target and it changes the selected object.
2. the option of manually editing the selected object without having to do heavy edits to a file (like removing multiple objects).

[krzysz00]

More nitpicking than actual issues - this translation seems reasonable

mlir/lib/Dialect/GPU/Targets/ObjectHandler.cpp

103 ↗

(On Diff #543193)

It feels weird to be declaring things into the llvm namespace from mlir/. Maybe this'd make sense in mlir::LLVM with a using?

I might also be off about how this translation code works

[fmorac]

mlir/lib/Dialect/GPU/Targets/ObjectHandler.cpp

103 ↗

(On Diff #543193)

The issue is conflicting names in namespaces, so I cannot have using mlir & using llvm, inside translation it made more sense to use llvm to fix this issue as it in its majority LLVM API.

[fmorac]

Moved the SelectObject implementation to the GPU translation library & added registration calls.

TODO:
The registration mechanism needs to be changed in the future, as the IR only verifies successfully if the registration call happens before verifying the IR. This issue also occurs with the NVVM & ROCDL registration mechanisms.
One solution is adding a promised flag in Tablegen that generates a separate trait without the interface methods and checks a promise was made. Also, the promised interface mechanism needs to allow finer granularity, allowing to specify specific attribute promises.

[mehdi_amini]

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

211

Isn't it an Attribute Interface? I'm confused what it means on the dialect here? Does it provide some fallback or something?

1654

cast?

[fmorac]

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

211

Yes, it's an attribute interface, however to not bundle LLVM libs into the GPUDialect lib and being able to put the implementation of the translation on the GPU Translation To LLVM lib, I have to register the interface as a promised interface.

As far as I could tell, the promised interface mechanism doesn't enable the granularity to say that something it's a promised interface for an specific attribute, it just promises an interface, that's why I listed it in the TODO on my previous comment:

The registration mechanism needs to be changed in the future, as the IR only verifies successfully if the registration call happens before verifying the IR. This issue also occurs with the NVVM & ROCDL registration mechanisms.
One solution is adding a promised flag in Tablegen that generates a separate trait without the interface methods and checks a promise was made. Also, the promised interface mechanism needs to allow finer granularity, allowing to specify specific attribute promises.

[mehdi_amini]

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

211

I don't understand what it means to promise this interface here actually?
What happens if we just remove this line without changing anything else?

[fmorac]

Remove the promised interface, instead use OffloadingTranslationAttrTrait in SelectObjectAttr
and register the interface automatically with the GPUTranslationInterface.

mlir/include/mlir/Dialect/GPU/IR/CompilationAttrs.td

def GPU_ObjectAttr : GPU_Attr<"Object", "object"> {
   }];
   let parameters = (ins "TargetAttrInterface":$target, "StringAttr":$object);
   let assemblyFormat = [{`<` $target `,` $object `>`}];
 }
 
 def GPUObjectArrayAttr :
   TypedArrayAttrBase<GPU_ObjectAttr, "an array of GPU object attributes">;
 
+//===----------------------------------------------------------------------===//
+// GPU offloading LLVM translation handler attributes.
+//===----------------------------------------------------------------------===//
+
+def GPU_SelectObjectAttr : GPU_Attr<"SelectObject", "select_object"> {
+  let description = [{
+    This GPU offloading handler selects a single GPU object for embedding. The
+    object is selected based on the `target` parameter, this parameter can be
+    either a number -i.e. selects the ith-target, or the target itself -i.e.
+    searches for the specified target in the object array.
+
+    The first object in a `gpu.binary` operation is selected if no target is
+    specified.
+  }];
+  let parameters = (ins
+    OptionalParameter<"Attribute", "Target to select for embedding.">:$target
+  );
+  let assemblyFormat = [{
+    (`<` $target^ `>`)?
+  }];
+  let genVerifyDecl = 1;
+}
+
 #endif // GPU_COMPILATION_ATTRS

mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h

public:
   using Base =
       typename Type::TypeBase<SparseHandleType<K>, Type, TypeStorage>::Base;
   using Base::Base;
 };
 
 using SparseDnTensorHandleType = SparseHandleType<SparseHandleKind::DnTensor>;
 using SparseSpMatHandleType = SparseHandleType<SparseHandleKind::SpMat>;
 
+/// Registers offloading LLVM translation interfaces. TODO: Remove this
+/// function.
+void registerOffloadingLLVMTranslationInterfacesExternalModels(
+    mlir::DialectRegistry &registry);
 } // namespace gpu
 } // namespace mlir
 
 #include "mlir/Dialect/GPU/IR/GPUOpsEnums.h.inc"
 
 #include "mlir/Dialect/GPU/IR/GPUOpsDialect.h.inc"
 
 #include "mlir/Dialect/GPU/IR/GPUOpInterfaces.h.inc"

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

let description = [{
 
     During translation into LLVM, the offloading attribute will be called
     for translating GPU binary and launch operations into LLVM instructions. If
     no attribute is provided, the default handler selects the first object from
     the array and embeds it as a string.
 
     Examples:
     ```
+      // Selects the first object.
       gpu.binary @myobject [#gpu.object<...>, #gpu.object<...>]
+      // Uses the `#foo.my_handler` for handling the binary during translation.
       gpu.binary @myobject <#foo.my_handler> [#gpu.object<...>, #gpu.object<...>]
+      // Selects the object with the `#rocdl.target` target attribute.
+      gpu.binary @myobject <#gpu.select_object<#rocdl.target>> [#gpu.object<...>, #gpu.object<#rocdl.target, ...>]
     ```
   }];
   let builders = [
     OpBuilder<(ins "StringRef":$name,
                    "OffloadingLLVMTranslationAttrInterface":$offloadingHandler,
                    "ArrayAttr":$objects)>,
     OpBuilder<(ins "StringRef":$name,
                    "OffloadingLLVMTranslationAttrInterface":$offloadingHandler,
                    "ArrayRef<Attribute>":$objects)>
   ];
   let skipDefaultBuilders = 1;
   let assemblyFormat = [{
-    $sym_name (`<` $offloadingHandler ^ `>`)? attr-dict $objects
+    $sym_name custom<OffloadingHandler>($offloadingHandler) attr-dict $objects
   }];
 }
 
 def GPU_HostRegisterOp : GPU_Op<"host_register">,
     Arguments<(ins AnyUnrankedMemRef:$value)> {
   let summary = "Registers a memref for access from device.";
   let description = [{
     This op maps the provided host buffer into the device address space.

mlir/include/mlir/InitAllDialects.h

inline void registerAllDialects(DialectRegistry &registry) {
 
   // Register all external models.
   affine::registerValueBoundsOpInterfaceExternalModels(registry);
   arith::registerBufferizableOpInterfaceExternalModels(registry);
   arith::registerValueBoundsOpInterfaceExternalModels(registry);
   bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(
       registry);
   builtin::registerCastOpInterfaceExternalModels(registry);
+  gpu::registerOffloadingLLVMTranslationInterfacesExternalModels(registry);
   linalg::registerBufferizableOpInterfaceExternalModels(registry);
   linalg::registerTilingInterfaceExternalModels(registry);
   linalg::registerValueBoundsOpInterfaceExternalModels(registry);
   memref::registerBufferizableOpInterfaceExternalModels(registry);
   memref::registerRuntimeVerifiableOpInterfaceExternalModels(registry);
   memref::registerValueBoundsOpInterfaceExternalModels(registry);
   memref::registerMemorySlotExternalModels(registry);
   scf::registerBufferizableOpInterfaceExternalModels(registry);

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

 #define GET_OP_LIST
 #include "mlir/Dialect/GPU/IR/GPUOps.cpp.inc"
       >();
   addAttributes<
 #define GET_ATTRDEF_LIST
 #include "mlir/Dialect/GPU/IR/GPUOpsAttributes.cpp.inc"
       >();
   addInterfaces<GPUInlinerInterface>();
+  declarePromisedInterface<OffloadingLLVMTranslationAttrInterface>();
 }
 
 static std::string getSparseHandleKeyword(SparseHandleKind kind) {
   switch (kind) {
   case SparseHandleKind::DnTensor:
     return "sparse.dntensor_handle";
   case SparseHandleKind::SpMat:
     return "sparse.spmat_handle";

 //===----------------------------------------------------------------------===//
 void BinaryOp::build(OpBuilder &builder, OperationState &result, StringRef name,
                      OffloadingLLVMTranslationAttrInterface offloadingHandler,
                      ArrayAttr objects) {
   auto &properties = result.getOrAddProperties<Properties>();
   result.attributes.push_back(builder.getNamedAttr(
       SymbolTable::getSymbolAttrName(), builder.getStringAttr(name)));
   properties.objects = objects;
+  if (offloadingHandler) {
     properties.offloadingHandler = offloadingHandler;
+  } else {
+    auto offloadingHandler = builder.getAttr<SelectObjectAttr>(nullptr);
+    properties.offloadingHandler =
+        dyn_cast<OffloadingLLVMTranslationAttrInterface>(offloadingHandler);
+  }
 }
 
 void BinaryOp::build(OpBuilder &builder, OperationState &result, StringRef name,
                      OffloadingLLVMTranslationAttrInterface offloadingHandler,
                      ArrayRef<Attribute> objects) {
   build(builder, result, name, offloadingHandler,
         objects.size() > 0 ? builder.getArrayAttr(objects) : ArrayAttr());
 }
 
+static ParseResult parseOffloadingHandler(OpAsmParser &parser,
+                                          Attribute &offloadingHandler) {
+  if (succeeded(parser.parseOptionalLess())) {
+    if (parser.parseAttribute(offloadingHandler))
+      return failure();
+    if (parser.parseGreater())
+      return failure();
+  }
+  if (!offloadingHandler)
+    offloadingHandler = parser.getBuilder().getAttr<SelectObjectAttr>(nullptr);
+  return success();
+}
+
+static void printOffloadingHandler(OpAsmPrinter &printer, Operation *op,
+                                   Attribute offloadingHandler) {
+  if (offloadingHandler != SelectObjectAttr::get(op->getContext(), nullptr))
+    printer << '<' << offloadingHandler << '>';
+}
+
 //===----------------------------------------------------------------------===//
 // GPUMemcpyOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult MemcpyOp::verify() {
   auto srcType = getSrc().getType();
   auto dstType = getDst().getType();
 

 } // namespace
 
 void AllocOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                           MLIRContext *context) {
   results.add<SimplifyDimOfAllocOp>(context);
 }
 
 //===----------------------------------------------------------------------===//
+// GPU select object attribute
+//===----------------------------------------------------------------------===//
+
+LogicalResult
+gpu::SelectObjectAttr::verify(function_ref<InFlightDiagnostic()> emitError,
+                              Attribute target) {
+  // Check `target`, it can be null, an integer attr or a GPU Target attribute.
+  if (target) {
+    if (auto intAttr = mlir::dyn_cast<IntegerAttr>(target)) {
+      if (intAttr.getInt() < 0) {
+        return emitError() << "The object index must be positive.";
+      }
+    } else if (!(::mlir::isa<TargetAttrInterface>(target))) {
+      return emitError()
+             << "The target attribute must be a GPU Target attribute.";
+    }
+  }
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // GPU target options
 //===----------------------------------------------------------------------===//
 
 TargetOptions::TargetOptions(StringRef toolkitPath,
                              ArrayRef<std::string> linkFiles,
                              StringRef cmdOptions,
                              CompilationTarget compilationTarget)
     : TargetOptions(TypeID::get<TargetOptions>(), toolkitPath, linkFiles,

mlir/lib/Target/LLVMIR/Dialect/GPU/CMakeLists.txt

 add_mlir_translation_library(MLIRGPUToLLVMIRTranslation
   GPUToLLVMIRTranslation.cpp
+  SelectObjectAttr.cpp
 
   LINK_COMPONENTS
   Core
 
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRGPUDialect
   MLIRLLVMDialect
   MLIRSupport
   MLIRTargetLLVMIRExport
   )

mlir/lib/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.cpp

 
 } // namespace
 
 void mlir::registerGPUDialectTranslation(DialectRegistry &registry) {
   registry.insert<gpu::GPUDialect>();
   registry.addExtension(+[](MLIRContext *ctx, gpu::GPUDialect *dialect) {
     dialect->addInterfaces<GPUDialectLLVMIRTranslationInterface>();
   });
+  gpu::registerOffloadingLLVMTranslationInterfacesExternalModels(registry);
 }
 
 void mlir::registerGPUDialectTranslation(MLIRContext &context) {
   DialectRegistry registry;
   registerGPUDialectTranslation(registry);
   context.appendDialectRegistry(registry);
 }

mlir/lib/Target/LLVMIR/Dialect/GPU/SelectObjectAttr.cpp

+//===- ObjectHandler.cpp - Implements base ObjectManager attributes -------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements the `OffloadingLLVMTranslationAttrInterface` for the
+// `SelectObject` attribute.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+
+#include "mlir/Target/LLVMIR/Export.h"
+#include "mlir/Target/LLVMIR/ModuleTranslation.h"
+
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Support/FormatVariadic.h"
+
+using namespace mlir;
+
+namespace {
+// Implementation of the `OffloadingLLVMTranslationAttrInterface` model.
+class SelectObjectAttrImpl
+    : public gpu::OffloadingLLVMTranslationAttrInterface::FallbackModel<
+          SelectObjectAttrImpl> {
+public:
+  // Translates a `gpu.binary`, embedding the binary into a host LLVM module as
+  // global binary string.
+  LogicalResult embedBinary(Attribute attribute, Operation *operation,
+                            llvm::IRBuilderBase &builder,
+                            LLVM::ModuleTranslation &moduleTranslation) const;
+
+  // Translates a `gpu.launch_func` to a sequence of LLVM instructions resulting
+  // in a kernel launch call.
+  LogicalResult launchKernel(Attribute attribute,
+                             Operation *launchFuncOperation,
+                             Operation *binaryOperation,
+                             llvm::IRBuilderBase &builder,
+                             LLVM::ModuleTranslation &moduleTranslation) const;
+};
+// Returns an identifier for the global string holding the binary.
+std::string getBinaryIdentifier(StringRef binaryName) {
+  return binaryName.str() + "_bin_cst";
+}
+} // namespace
+
+void mlir::gpu::registerOffloadingLLVMTranslationInterfacesExternalModels(
+    DialectRegistry &registry) {
+  registry.addExtension(+[](MLIRContext *ctx, gpu::GPUDialect *dialect) {
+    SelectObjectAttr::attachInterface<SelectObjectAttrImpl>(*ctx);
+  });
+}
+
+LogicalResult SelectObjectAttrImpl::embedBinary(
+    Attribute attribute, Operation *operation, llvm::IRBuilderBase &builder,
+    LLVM::ModuleTranslation &moduleTranslation) const {
+  assert(operation && "The binary operation must be non null.");
+  if (!operation)
+    return failure();
+
+  auto op = mlir::dyn_cast<gpu::BinaryOp>(operation);
+  if (!op) {
+    operation->emitError("Operation must be a GPU binary.");
+    return failure();
+  }
+
+  ArrayRef<Attribute> objects = op.getObjectsAttr().getValue();
+
+  // Obtain the index of the object to select.
+  int64_t index = -1;
+  if (Attribute target = cast<gpu::SelectObjectAttr>(attribute).getTarget()) {
+    // If the target attribute is a number it is the index. Otherwise compare
+    // the attribute to every target inside the object array to find the index.
+    if (auto indexAttr = mlir::dyn_cast<IntegerAttr>(target)) {
+      index = indexAttr.getInt();
+    } else {
+      for (auto [i, attr] : llvm::enumerate(objects)) {
+        auto obj = mlir::dyn_cast<gpu::ObjectAttr>(attr);
+        if (obj.getTarget() == target) {
+          index = i;
+        }
+      }
+    }
+  } else {
+    // If the target attribute is null then it's selecting the first object in
+    // the object array.
+    index = 0;
+  }
+
+  if (index < 0 || index >= static_cast<int64_t>(objects.size())) {
+    op->emitError("The requested target object couldn't be found.");
+    return failure();
+  }
+  auto object = mlir::dyn_cast<gpu::ObjectAttr>(objects[index]);
+
+  llvm::Module *module = moduleTranslation.getLLVMModule();
+
+  // Embed the object as a global string.
+  llvm::Constant *binary = llvm::ConstantDataArray::getString(
+      builder.getContext(), object.getObject().getValue(), false);
+  llvm::GlobalVariable *serializedObj =
+      new llvm::GlobalVariable(*module, binary->getType(), true,
+                               llvm::GlobalValue::LinkageTypes::InternalLinkage,
+                               binary, getBinaryIdentifier(op.getName()));
+  serializedObj->setLinkage(llvm::GlobalValue::LinkageTypes::InternalLinkage);
+  serializedObj->setAlignment(llvm::MaybeAlign(8));
+  serializedObj->setUnnamedAddr(llvm::GlobalValue::UnnamedAddr::None);
+  return success();
+}
+
+namespace llvm {
+namespace {
+class LaunchKernel {
+public:
+  LaunchKernel(Module &module, IRBuilderBase &builder,
+               mlir::LLVM::ModuleTranslation &moduleTranslation);
+  // Get the kernel launch callee.
+  FunctionCallee getKernelLaunchFn();
+
+  // Get the module function callee.
+  FunctionCallee getModuleFunctionFn();
+
+  // Get the module load callee.
+  FunctionCallee getModuleLoadFn();
+
+  // Get the module unload callee.
+  FunctionCallee getModuleUnloadFn();
+
+  // Get the stream create callee.
+  FunctionCallee getStreamCreateFn();
+
+  // Get the stream destroy callee.
+  FunctionCallee getStreamDestroyFn();
+
+  // Get the stream sync callee.
+  FunctionCallee getStreamSyncFn();
+
+  // Ger or create the function name global string.
+  Value *getOrCreateFunctionName(StringRef moduleName, StringRef kernelName);
+
+  // Create the void* kernel array for passing the arguments.
+  Value *createKernelArgArray(mlir::gpu::LaunchFuncOp op);
+
+  // Create the full kernel launch.
+  mlir::LogicalResult createKernelLaunch(mlir::gpu::LaunchFuncOp op);
+
+private:
+  Module &module;
+  IRBuilderBase &builder;
+  mlir::LLVM::ModuleTranslation &moduleTranslation;
+  Type *i32Ty{};
+  Type *voidTy{};
+  Type *intPtrTy{};
+  PointerType *ptrTy{};
+};
+} // namespace
+} // namespace llvm
+
+LogicalResult SelectObjectAttrImpl::launchKernel(
+    Attribute attribute, Operation *launchFuncOperation,
+    Operation *binaryOperation, llvm::IRBuilderBase &builder,
+    LLVM::ModuleTranslation &moduleTranslation) const {
+
+  assert(launchFuncOperation && "The launch func operation must be non null.");
+  if (!launchFuncOperation)
+    return failure();
+
+  auto launchFuncOp = mlir::dyn_cast<gpu::LaunchFuncOp>(launchFuncOperation);
+  if (!launchFuncOp) {
+    launchFuncOperation->emitError("Operation must be a GPU launch func Op.");
+    return failure();
+  }
+
+  return llvm::LaunchKernel(*moduleTranslation.getLLVMModule(), builder,
+                            moduleTranslation)
+      .createKernelLaunch(launchFuncOp);
+}
+
+llvm::LaunchKernel::LaunchKernel(
+    Module &module, IRBuilderBase &builder,
+    mlir::LLVM::ModuleTranslation &moduleTranslation)
+    : module(module), builder(builder), moduleTranslation(moduleTranslation) {
+  i32Ty = builder.getInt32Ty();
+  ptrTy = builder.getPtrTy(0);
+  voidTy = builder.getVoidTy();
+  intPtrTy = builder.getIntPtrTy(module.getDataLayout());
+}
+
+llvm::FunctionCallee llvm::LaunchKernel::getKernelLaunchFn() {
+  return module.getOrInsertFunction(
+      "mgpuLaunchKernel",
+      FunctionType::get(
+          voidTy,
+          ArrayRef<Type *>({ptrTy, intPtrTy, intPtrTy, intPtrTy, intPtrTy,
+                            intPtrTy, intPtrTy, i32Ty, ptrTy, ptrTy, ptrTy}),
+          false));
+}
+
+llvm::FunctionCallee llvm::LaunchKernel::getModuleFunctionFn() {
+  return module.getOrInsertFunction(
+      "mgpuModuleGetFunction",
+      FunctionType::get(ptrTy, ArrayRef<Type *>({ptrTy, ptrTy}), false));
+}
+
+llvm::FunctionCallee llvm::LaunchKernel::getModuleLoadFn() {
+  return module.getOrInsertFunction(
+      "mgpuModuleLoad",
+      FunctionType::get(ptrTy, ArrayRef<Type *>({ptrTy}), false));
+}
+
+llvm::FunctionCallee llvm::LaunchKernel::getModuleUnloadFn() {
+  return module.getOrInsertFunction(
+      "mgpuModuleUnload",
+      FunctionType::get(voidTy, ArrayRef<Type *>({ptrTy}), false));
+}
+
+llvm::FunctionCallee llvm::LaunchKernel::getStreamCreateFn() {
+  return module.getOrInsertFunction("mgpuStreamCreate",
+                                    FunctionType::get(ptrTy, false));
+}
+
+llvm::FunctionCallee llvm::LaunchKernel::getStreamDestroyFn() {
+  return module.getOrInsertFunction(
+      "mgpuStreamDestroy",
+      FunctionType::get(voidTy, ArrayRef<Type *>({ptrTy}), false));
+}
+
+llvm::FunctionCallee llvm::LaunchKernel::getStreamSyncFn() {
+  return module.getOrInsertFunction(
+      "mgpuStreamSynchronize",
+      FunctionType::get(voidTy, ArrayRef<Type *>({ptrTy}), false));
+}
+
+// Generates an LLVM IR dialect global that contains the name of the given
+// kernel function as a C string, and returns a pointer to its beginning.
+llvm::Value *llvm::LaunchKernel::getOrCreateFunctionName(StringRef moduleName,
+                                                         StringRef kernelName) {
+  std::string globalName =
+      std::string(formatv("{0}_{1}_kernel_name", moduleName, kernelName));
+
+  if (GlobalVariable *gv = module.getGlobalVariable(globalName))
+    return gv;
+
+  return builder.CreateGlobalString(kernelName, globalName);
+}
+
+// Creates a struct containing all kernel parameters on the stack and returns
+// an array of type-erased pointers to the fields of the struct. The array can
+// then be passed to the CUDA / ROCm (HIP) kernel launch calls.
+// The generated code is essentially as follows:
+//
+// %struct = alloca(sizeof(struct { Parameters... }))
+// %array = alloca(NumParameters * sizeof(void *))
+// for (i : [0, NumParameters))
+//   %fieldPtr = llvm.getelementptr %struct[0, i]
+//   llvm.store parameters[i], %fieldPtr
+//   %elementPtr = llvm.getelementptr %array[i]
+//   llvm.store %fieldPtr, %elementPtr
+// return %array
+llvm::Value *
+llvm::LaunchKernel::createKernelArgArray(mlir::gpu::LaunchFuncOp op) {
+  SmallVector<Value *> args =
+      moduleTranslation.lookupValues(op.getKernelOperands());
+  SmallVector<Type *> structTypes(args.size(), nullptr);
+
+  for (auto [i, arg] : llvm::enumerate(args))
+    structTypes[i] = arg->getType();
+
+  Type *structTy = StructType::create(module.getContext(), structTypes);
+  Value *argStruct = builder.CreateAlloca(structTy, 0u);
+  Value *argArray = builder.CreateAlloca(
+      ptrTy, ConstantInt::get(intPtrTy, structTypes.size()));
+
+  for (auto [i, arg] : enumerate(args)) {
+    Value *structMember = builder.CreateStructGEP(structTy, argStruct, i);
+    builder.CreateStore(arg, structMember);
+    Value *arrayMember = builder.CreateConstGEP1_32(ptrTy, argArray, i);
+    builder.CreateStore(structMember, arrayMember);
+  }
+  return argArray;
+}
+
+// Emits LLVM IR to launch a kernel function:
+// %0 = call %binarygetter
+// %1 = call %moduleLoad(%0)
+// %2 = <see generateKernelNameConstant>
+// %3 = call %moduleGetFunction(%1, %2)
+// %4 = call %streamCreate()
+// %5 = <see generateParamsArray>
+// call %launchKernel(%3, <launchOp operands 0..5>, 0, %4, %5, nullptr)
+// call %streamSynchronize(%4)
+// call %streamDestroy(%4)
+// call %moduleUnload(%1)
+mlir::LogicalResult
+llvm::LaunchKernel::createKernelLaunch(mlir::gpu::LaunchFuncOp op) {
+  auto llvmValue = [&](mlir::Value value) -> Value * {
+    Value *v = moduleTranslation.lookupValue(value);
+    assert(v && "Value has not been translated.");
+    return v;
+  };
+
+  // Get grid dimensions.
+  mlir::gpu::KernelDim3 grid = op.getGridSizeOperandValues();
+  Value *gx = llvmValue(grid.x), *gy = llvmValue(grid.y),
+        *gz = llvmValue(grid.z);
+
+  // Get block dimensions.
+  mlir::gpu::KernelDim3 block = op.getBlockSizeOperandValues();
+  Value *bx = llvmValue(block.x), *by = llvmValue(block.y),
+        *bz = llvmValue(block.z);
+
+  // Get dynamic shared memory size.
+  Value *dynamicMemorySize = nullptr;
+  if (mlir::Value dynSz = op.getDynamicSharedMemorySize())
+    dynamicMemorySize = llvmValue(dynSz);
+  else
+    dynamicMemorySize = ConstantInt::get(i32Ty, 0);
+
+  // Create the argument array.
+  Value *argArray = createKernelArgArray(op);
+
+  // Load the kernel module.
+  StringRef moduleName = op.getKernelModuleName().getValue();
+  std::string binaryIdentifier = getBinaryIdentifier(moduleName);
+  Value *binary = module.getGlobalVariable(binaryIdentifier, true);
+  if (!binary)
+    return op.emitError() << "Couldn't find the binary: " << binaryIdentifier;
+  Value *moduleObject = builder.CreateCall(getModuleLoadFn(), {binary});
+
+  // Load the kernel function.
+  Value *moduleFunction = builder.CreateCall(
+      getModuleFunctionFn(),
+      {moduleObject,
+       getOrCreateFunctionName(moduleName, op.getKernelName().getValue())});
+
+  // Get the stream to use for execution. If there's no async object then create
+  // a stream to make a synchronous kernel launch.
+  Value *stream = nullptr;
+  bool handleStream = false;
+  if (mlir::Value asyncObject = op.getAsyncObject()) {
+    stream = llvmValue(asyncObject);
+  } else {
+    handleStream = true;
+    stream = builder.CreateCall(getStreamCreateFn(), {});
+  }
+
+  // Create the launch call.
+  Value *nullPtr = ConstantPointerNull::get(ptrTy);
+  builder.CreateCall(
+      getKernelLaunchFn(),
+      ArrayRef<Value *>({moduleFunction, gx, gy, gz, bx, by, bz,
+                         dynamicMemorySize, stream, argArray, nullPtr}));
+
+  // Sync & destroy the stream, for synchronous launches.
+  if (handleStream) {
+    builder.CreateCall(getStreamSyncFn(), {stream});
+    builder.CreateCall(getStreamDestroyFn(), {stream});
+  }
+
+  // Unload the kernel module.
+  builder.CreateCall(getModuleUnloadFn(), {moduleObject});
+
+  return success();
+}

mlir/test/Dialect/GPU/invalid.mlir

 module {
   // expected-error @+1 {{'gpu.binary' op attribute 'objects' failed to satisfy constraint: an array of GPU object attributes with at least 1 elements}}
   gpu.binary @binary []
 }
 
 // -----
 
 module {
-  // expected-error @+1 {{custom op 'gpu.binary' invalid kind of attribute specified}}
+  // expected-error @+1 {{'gpu.binary' op attribute 'offloadingHandler' failed to satisfy constraint: OffloadingLLVMTranslationAttrInterface instance}}
   gpu.binary @binary <1> [#gpu.object<#nvvm.target, "">]
 }

mlir/test/Dialect/GPU/ops.mlir

gpu.func @kernel_1(%arg0 : f32, %arg1 : memref<?xf32, 1>) kernel {
       gpu.return
     }
 
     gpu.func @kernel_2() kernel {
       gpu.return
     }
   }
 
+  gpu.binary @binary_1 [#gpu.object<#nvvm.target, "">]
+
+  gpu.binary @binary_2 <#gpu.select_object<#rocdl.target>> [#gpu.object<#nvvm.target, "">, #gpu.object<#rocdl.target, "">]
+
+  gpu.binary @binary_3 <#gpu.select_object<1>> [#gpu.object<#nvvm.target, "">, #gpu.object<#rocdl.target, "">]
+
   func.func private @two_value_generator() -> (f32, memref<?xf32, 1>)
 
   func.func @foo() {
     %0 = "op"() : () -> (f32)
     %1 = "op"() : () -> (memref<?xf32, 1>)
     // CHECK: %{{.*}} = arith.constant 8
     %cst = arith.constant 8 : index
     %cstI64 = arith.constant 8 : i64
func.func @foo() {
     %t1 = gpu.launch_func async [%t0] @kernels::@kernel_2  blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst)
 
     // CHECK: gpu.launch_func <%{{.*}} : !llvm.ptr> @kernels::@kernel_1 blocks in (%{{.*}}, %{{.*}}, %{{.*}}) threads in (%{{.*}}, %{{.*}}, %{{.*}}) : i64 args(%{{.*}} : f32, %{{.*}} : memref<?xf32, 1>)
     gpu.launch_func <%lowStream : !llvm.ptr> @kernels::@kernel_1 blocks in (%cstI64, %cstI64, %cstI64) threads in (%cstI64, %cstI64, %cstI64) : i64 args(%0 : f32, %1 : memref<?xf32, 1>)
 
     // CHECK: gpu.launch_func @kernels::@kernel_1 blocks in (%{{.*}}, %{{.*}}, %{{.*}}) threads in (%{{.*}}, %{{.*}}, %{{.*}}) : i32 args(%{{.*}} : f32, %{{.*}} : memref<?xf32, 1>)
     gpu.launch_func @kernels::@kernel_1 blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0) : i32 args(%0 : f32, %1 : memref<?xf32, 1>)
 
+    // CHECK: gpu.launch_func @binary_1::@kernel blocks in (%{{.*}}, %{{.*}}, %{{.*}}) threads in (%{{.*}}, %{{.*}}, %{{.*}}) : i32 args(%{{.*}} : f32, %{{.*}} : memref<?xf32, 1>)
+    gpu.launch_func @binary_1::@kernel blocks in (%c0, %c0, %c0) threads in (%c0, %c0, %c0) : i32 args(%0 : f32, %1 : memref<?xf32, 1>)
+
     // CHECK: %[[VALUES:.*]]:2 = call
     %values:2 = func.call @two_value_generator() : () -> (f32, memref<?xf32, 1>)
     // CHECK: gpu.launch_func @kernels::@kernel_1 {{.*}} args(%[[VALUES]]#0 : f32, %[[VALUES]]#1 : memref<?xf32, 1>)
     gpu.launch_func @kernels::@kernel_1 blocks in (%cst, %cst, %cst) threads in (%cst, %cst, %cst) args(%values#0 : f32, %values#1 : memref<?xf32, 1>)
 
     return
   }