MLIR  18.0.0git
SelectObjectAttr.cpp
Go to the documentation of this file.
1 //===- ObjectHandler.cpp - Implements base ObjectManager attributes -------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file implements the `OffloadingLLVMTranslationAttrInterface` for the
10 // `SelectObject` attribute.
11 //
12 //===----------------------------------------------------------------------===//
13 
15 
19 
20 #include "llvm/IR/Constants.h"
21 #include "llvm/IR/IRBuilder.h"
22 #include "llvm/IR/LLVMContext.h"
23 #include "llvm/IR/Module.h"
24 #include "llvm/Support/FormatVariadic.h"
25 
26 using namespace mlir;
27 
28 namespace {
29 // Implementation of the `OffloadingLLVMTranslationAttrInterface` model.
30 class SelectObjectAttrImpl
31  : public gpu::OffloadingLLVMTranslationAttrInterface::FallbackModel<
32  SelectObjectAttrImpl> {
33 public:
34  // Translates a `gpu.binary`, embedding the binary into a host LLVM module as
35  // global binary string.
36  LogicalResult embedBinary(Attribute attribute, Operation *operation,
37  llvm::IRBuilderBase &builder,
38  LLVM::ModuleTranslation &moduleTranslation) const;
39 
40  // Translates a `gpu.launch_func` to a sequence of LLVM instructions resulting
41  // in a kernel launch call.
43  Operation *launchFuncOperation,
44  Operation *binaryOperation,
45  llvm::IRBuilderBase &builder,
46  LLVM::ModuleTranslation &moduleTranslation) const;
47 
48  // Returns the selected object for embedding.
49  gpu::ObjectAttr getSelectedObject(gpu::BinaryOp op) const;
50 };
51 // Returns an identifier for the global string holding the binary.
52 std::string getBinaryIdentifier(StringRef binaryName) {
53  return binaryName.str() + "_bin_cst";
54 }
55 } // namespace
56 
58  DialectRegistry &registry) {
59  registry.addExtension(+[](MLIRContext *ctx, gpu::GPUDialect *dialect) {
60  SelectObjectAttr::attachInterface<SelectObjectAttrImpl>(*ctx);
61  });
62 }
63 
64 gpu::ObjectAttr
65 SelectObjectAttrImpl::getSelectedObject(gpu::BinaryOp op) const {
66  ArrayRef<Attribute> objects = op.getObjectsAttr().getValue();
67 
68  // Obtain the index of the object to select.
69  int64_t index = -1;
70  if (Attribute target =
71  cast<gpu::SelectObjectAttr>(op.getOffloadingHandlerAttr())
72  .getTarget()) {
73  // If the target attribute is a number it is the index. Otherwise compare
74  // the attribute to every target inside the object array to find the index.
75  if (auto indexAttr = mlir::dyn_cast<IntegerAttr>(target)) {
76  index = indexAttr.getInt();
77  } else {
78  for (auto [i, attr] : llvm::enumerate(objects)) {
79  auto obj = mlir::dyn_cast<gpu::ObjectAttr>(attr);
80  if (obj.getTarget() == target) {
81  index = i;
82  }
83  }
84  }
85  } else {
86  // If the target attribute is null then it's selecting the first object in
87  // the object array.
88  index = 0;
89  }
90 
91  if (index < 0 || index >= static_cast<int64_t>(objects.size())) {
92  op->emitError("the requested target object couldn't be found");
93  return nullptr;
94  }
95  return mlir::dyn_cast<gpu::ObjectAttr>(objects[index]);
96 }
97 
98 LogicalResult SelectObjectAttrImpl::embedBinary(
99  Attribute attribute, Operation *operation, llvm::IRBuilderBase &builder,
100  LLVM::ModuleTranslation &moduleTranslation) const {
101  assert(operation && "The binary operation must be non null.");
102  if (!operation)
103  return failure();
104 
105  auto op = mlir::dyn_cast<gpu::BinaryOp>(operation);
106  if (!op) {
107  operation->emitError("operation must be a GPU binary");
108  return failure();
109  }
110 
111  gpu::ObjectAttr object = getSelectedObject(op);
112  if (!object)
113  return failure();
114 
115  llvm::Module *module = moduleTranslation.getLLVMModule();
116 
117  // Embed the object as a global string.
118  llvm::Constant *binary = llvm::ConstantDataArray::getString(
119  builder.getContext(), object.getObject().getValue(), false);
120  llvm::GlobalVariable *serializedObj =
121  new llvm::GlobalVariable(*module, binary->getType(), true,
122  llvm::GlobalValue::LinkageTypes::InternalLinkage,
123  binary, getBinaryIdentifier(op.getName()));
124  serializedObj->setLinkage(llvm::GlobalValue::LinkageTypes::InternalLinkage);
125  serializedObj->setAlignment(llvm::MaybeAlign(8));
126  serializedObj->setUnnamedAddr(llvm::GlobalValue::UnnamedAddr::None);
127  return success();
128 }
129 
130 namespace llvm {
131 namespace {
132 class LaunchKernel {
133 public:
134  LaunchKernel(Module &module, IRBuilderBase &builder,
135  mlir::LLVM::ModuleTranslation &moduleTranslation);
136  // Get the kernel launch callee.
137  FunctionCallee getKernelLaunchFn();
138 
139  // Get the kernel launch callee.
140  FunctionCallee getClusterKernelLaunchFn();
141 
142  // Get the module function callee.
143  FunctionCallee getModuleFunctionFn();
144 
145  // Get the module load callee.
146  FunctionCallee getModuleLoadFn();
147 
148  // Get the module load JIT callee.
149  FunctionCallee getModuleLoadJITFn();
150 
151  // Get the module unload callee.
152  FunctionCallee getModuleUnloadFn();
153 
154  // Get the stream create callee.
155  FunctionCallee getStreamCreateFn();
156 
157  // Get the stream destroy callee.
158  FunctionCallee getStreamDestroyFn();
159 
160  // Get the stream sync callee.
161  FunctionCallee getStreamSyncFn();
162 
163  // Ger or create the function name global string.
164  Value *getOrCreateFunctionName(StringRef moduleName, StringRef kernelName);
165 
166  // Create the void* kernel array for passing the arguments.
167  Value *createKernelArgArray(mlir::gpu::LaunchFuncOp op);
168 
169  // Create the full kernel launch.
170  mlir::LogicalResult createKernelLaunch(mlir::gpu::LaunchFuncOp op,
171  mlir::gpu::ObjectAttr object);
172 
173 private:
174  Module &module;
175  IRBuilderBase &builder;
176  mlir::LLVM::ModuleTranslation &moduleTranslation;
177  Type *i32Ty{};
178  Type *voidTy{};
179  Type *intPtrTy{};
180  PointerType *ptrTy{};
181 };
182 } // namespace
183 } // namespace llvm
184 
186  Attribute attribute, Operation *launchFuncOperation,
187  Operation *binaryOperation, llvm::IRBuilderBase &builder,
188  LLVM::ModuleTranslation &moduleTranslation) const {
189 
190  assert(launchFuncOperation && "The launch func operation must be non null.");
191  if (!launchFuncOperation)
192  return failure();
193 
194  auto launchFuncOp = mlir::dyn_cast<gpu::LaunchFuncOp>(launchFuncOperation);
195  if (!launchFuncOp) {
196  launchFuncOperation->emitError("operation must be a GPU launch func Op.");
197  return failure();
198  }
199 
200  auto binOp = mlir::dyn_cast<gpu::BinaryOp>(binaryOperation);
201  if (!binOp) {
202  binaryOperation->emitError("operation must be a GPU binary.");
203  return failure();
204  }
205  gpu::ObjectAttr object = getSelectedObject(binOp);
206  if (!object)
207  return failure();
208 
209  return llvm::LaunchKernel(*moduleTranslation.getLLVMModule(), builder,
210  moduleTranslation)
211  .createKernelLaunch(launchFuncOp, object);
212 }
213 
214 llvm::LaunchKernel::LaunchKernel(
215  Module &module, IRBuilderBase &builder,
216  mlir::LLVM::ModuleTranslation &moduleTranslation)
217  : module(module), builder(builder), moduleTranslation(moduleTranslation) {
218  i32Ty = builder.getInt32Ty();
219  ptrTy = builder.getPtrTy(0);
220  voidTy = builder.getVoidTy();
221  intPtrTy = builder.getIntPtrTy(module.getDataLayout());
222 }
223 
224 llvm::FunctionCallee llvm::LaunchKernel::getKernelLaunchFn() {
225  return module.getOrInsertFunction(
226  "mgpuLaunchKernel",
228  voidTy,
229  ArrayRef<Type *>({ptrTy, intPtrTy, intPtrTy, intPtrTy, intPtrTy,
230  intPtrTy, intPtrTy, i32Ty, ptrTy, ptrTy, ptrTy}),
231  false));
232 }
233 
234 llvm::FunctionCallee llvm::LaunchKernel::getClusterKernelLaunchFn() {
235  return module.getOrInsertFunction(
236  "mgpuLaunchClusterKernel",
238  voidTy,
239  ArrayRef<Type *>({ptrTy, intPtrTy, intPtrTy, intPtrTy, intPtrTy,
240  intPtrTy, intPtrTy, intPtrTy, intPtrTy, intPtrTy,
241  i32Ty, ptrTy, ptrTy, ptrTy}),
242  false));
243 }
244 
245 llvm::FunctionCallee llvm::LaunchKernel::getModuleFunctionFn() {
246  return module.getOrInsertFunction(
247  "mgpuModuleGetFunction",
248  FunctionType::get(ptrTy, ArrayRef<Type *>({ptrTy, ptrTy}), false));
249 }
250 
251 llvm::FunctionCallee llvm::LaunchKernel::getModuleLoadFn() {
252  return module.getOrInsertFunction(
253  "mgpuModuleLoad",
254  FunctionType::get(ptrTy, ArrayRef<Type *>({ptrTy}), false));
255 }
256 
257 llvm::FunctionCallee llvm::LaunchKernel::getModuleLoadJITFn() {
258  return module.getOrInsertFunction(
259  "mgpuModuleLoadJIT",
260  FunctionType::get(ptrTy, ArrayRef<Type *>({ptrTy, i32Ty}), false));
261 }
262 
263 llvm::FunctionCallee llvm::LaunchKernel::getModuleUnloadFn() {
264  return module.getOrInsertFunction(
265  "mgpuModuleUnload",
266  FunctionType::get(voidTy, ArrayRef<Type *>({ptrTy}), false));
267 }
268 
269 llvm::FunctionCallee llvm::LaunchKernel::getStreamCreateFn() {
270  return module.getOrInsertFunction("mgpuStreamCreate",
271  FunctionType::get(ptrTy, false));
272 }
273 
274 llvm::FunctionCallee llvm::LaunchKernel::getStreamDestroyFn() {
275  return module.getOrInsertFunction(
276  "mgpuStreamDestroy",
277  FunctionType::get(voidTy, ArrayRef<Type *>({ptrTy}), false));
278 }
279 
280 llvm::FunctionCallee llvm::LaunchKernel::getStreamSyncFn() {
281  return module.getOrInsertFunction(
282  "mgpuStreamSynchronize",
283  FunctionType::get(voidTy, ArrayRef<Type *>({ptrTy}), false));
284 }
285 
286 // Generates an LLVM IR dialect global that contains the name of the given
287 // kernel function as a C string, and returns a pointer to its beginning.
288 llvm::Value *llvm::LaunchKernel::getOrCreateFunctionName(StringRef moduleName,
289  StringRef kernelName) {
290  std::string globalName =
291  std::string(formatv("{0}_{1}_kernel_name", moduleName, kernelName));
292 
293  if (GlobalVariable *gv = module.getGlobalVariable(globalName))
294  return gv;
295 
296  return builder.CreateGlobalString(kernelName, globalName);
297 }
298 
299 // Creates a struct containing all kernel parameters on the stack and returns
300 // an array of type-erased pointers to the fields of the struct. The array can
301 // then be passed to the CUDA / ROCm (HIP) kernel launch calls.
302 // The generated code is essentially as follows:
303 //
304 // %struct = alloca(sizeof(struct { Parameters... }))
305 // %array = alloca(NumParameters * sizeof(void *))
306 // for (i : [0, NumParameters))
307 // %fieldPtr = llvm.getelementptr %struct[0, i]
308 // llvm.store parameters[i], %fieldPtr
309 // %elementPtr = llvm.getelementptr %array[i]
310 // llvm.store %fieldPtr, %elementPtr
311 // return %array
312 llvm::Value *
313 llvm::LaunchKernel::createKernelArgArray(mlir::gpu::LaunchFuncOp op) {
314  SmallVector<Value *> args =
315  moduleTranslation.lookupValues(op.getKernelOperands());
316  SmallVector<Type *> structTypes(args.size(), nullptr);
317 
318  for (auto [i, arg] : llvm::enumerate(args))
319  structTypes[i] = arg->getType();
320 
321  Type *structTy = StructType::create(module.getContext(), structTypes);
322  Value *argStruct = builder.CreateAlloca(structTy, 0u);
323  Value *argArray = builder.CreateAlloca(
324  ptrTy, ConstantInt::get(intPtrTy, structTypes.size()));
325 
326  for (auto [i, arg] : enumerate(args)) {
327  Value *structMember = builder.CreateStructGEP(structTy, argStruct, i);
328  builder.CreateStore(arg, structMember);
329  Value *arrayMember = builder.CreateConstGEP1_32(ptrTy, argArray, i);
330  builder.CreateStore(structMember, arrayMember);
331  }
332  return argArray;
333 }
334 
335 // Emits LLVM IR to launch a kernel function:
336 // %0 = call %binarygetter
337 // %1 = call %moduleLoad(%0)
338 // %2 = <see generateKernelNameConstant>
339 // %3 = call %moduleGetFunction(%1, %2)
340 // %4 = call %streamCreate()
341 // %5 = <see generateParamsArray>
342 // call %launchKernel(%3, <launchOp operands 0..5>, 0, %4, %5, nullptr)
343 // call %streamSynchronize(%4)
344 // call %streamDestroy(%4)
345 // call %moduleUnload(%1)
347 llvm::LaunchKernel::createKernelLaunch(mlir::gpu::LaunchFuncOp op,
348  mlir::gpu::ObjectAttr object) {
349  auto llvmValue = [&](mlir::Value value) -> Value * {
350  Value *v = moduleTranslation.lookupValue(value);
351  assert(v && "Value has not been translated.");
352  return v;
353  };
354 
355  // Get grid dimensions.
356  mlir::gpu::KernelDim3 grid = op.getGridSizeOperandValues();
357  Value *gx = llvmValue(grid.x), *gy = llvmValue(grid.y),
358  *gz = llvmValue(grid.z);
359 
360  // Get block dimensions.
361  mlir::gpu::KernelDim3 block = op.getBlockSizeOperandValues();
362  Value *bx = llvmValue(block.x), *by = llvmValue(block.y),
363  *bz = llvmValue(block.z);
364 
365  // Get dynamic shared memory size.
366  Value *dynamicMemorySize = nullptr;
367  if (mlir::Value dynSz = op.getDynamicSharedMemorySize())
368  dynamicMemorySize = llvmValue(dynSz);
369  else
370  dynamicMemorySize = ConstantInt::get(i32Ty, 0);
371 
372  // Create the argument array.
373  Value *argArray = createKernelArgArray(op);
374 
375  // Default JIT optimization level.
376  llvm::Constant *optV = llvm::ConstantInt::get(i32Ty, 0);
377  // Check if there's an optimization level embedded in the object.
378  DictionaryAttr objectProps = object.getProperties();
379  mlir::Attribute optAttr;
380  if (objectProps && (optAttr = objectProps.get("O"))) {
381  auto optLevel = dyn_cast<IntegerAttr>(optAttr);
382  if (!optLevel)
383  return op.emitError("the optimization level must be an integer");
384  optV = llvm::ConstantInt::get(i32Ty, optLevel.getValue());
385  }
386 
387  // Load the kernel module.
388  StringRef moduleName = op.getKernelModuleName().getValue();
389  std::string binaryIdentifier = getBinaryIdentifier(moduleName);
390  Value *binary = module.getGlobalVariable(binaryIdentifier, true);
391  if (!binary)
392  return op.emitError() << "Couldn't find the binary: " << binaryIdentifier;
393 
394  Value *moduleObject =
395  object.getFormat() == gpu::CompilationTarget::Assembly
396  ? builder.CreateCall(getModuleLoadJITFn(), {binary, optV})
397  : builder.CreateCall(getModuleLoadFn(), {binary});
398 
399  // Load the kernel function.
400  Value *moduleFunction = builder.CreateCall(
401  getModuleFunctionFn(),
402  {moduleObject,
403  getOrCreateFunctionName(moduleName, op.getKernelName().getValue())});
404 
405  // Get the stream to use for execution. If there's no async object then create
406  // a stream to make a synchronous kernel launch.
407  Value *stream = nullptr;
408  bool handleStream = false;
409  if (mlir::Value asyncObject = op.getAsyncObject()) {
410  stream = llvmValue(asyncObject);
411  } else {
412  handleStream = true;
413  stream = builder.CreateCall(getStreamCreateFn(), {});
414  }
415 
416  // Create the launch call.
417  Value *nullPtr = ConstantPointerNull::get(ptrTy);
418 
419  // Launch kernel with clusters if cluster size is specified.
420  if (op.hasClusterSize()) {
421  mlir::gpu::KernelDim3 cluster = op.getClusterSizeOperandValues();
422  Value *cx = llvmValue(cluster.x), *cy = llvmValue(cluster.y),
423  *cz = llvmValue(cluster.z);
424  builder.CreateCall(
425  getClusterKernelLaunchFn(),
426  ArrayRef<Value *>({moduleFunction, cx, cy, cz, gx, gy, gz, bx, by, bz,
427  dynamicMemorySize, stream, argArray, nullPtr}));
428  } else {
429  builder.CreateCall(
430  getKernelLaunchFn(),
431  ArrayRef<Value *>({moduleFunction, gx, gy, gz, bx, by, bz,
432  dynamicMemorySize, stream, argArray, nullPtr}));
433  }
434 
435  // Sync & destroy the stream, for synchronous launches.
436  if (handleStream) {
437  builder.CreateCall(getStreamSyncFn(), {stream});
438  builder.CreateCall(getStreamDestroyFn(), {stream});
439  }
440 
441  // Unload the kernel module.
442  builder.CreateCall(getModuleUnloadFn(), {moduleObject});
443 
444  return success();
445 }
@ None
static void launchKernel(sycl::queue *queue, sycl::kernel *kernel, size_t gridX, size_t gridY, size_t gridZ, size_t blockX, size_t blockY, size_t blockZ, size_t sharedMemBytes, void **params, size_t paramsCount)
Attributes are known-constant values of operations.
Definition: Attributes.h:25
The DialectRegistry maps a dialect namespace to a constructor for the matching dialect.
void addExtension(std::unique_ptr< DialectExtensionBase > extension)
Add the given extension to the registry.
Implementation class for module translation.
llvm::Value * lookupValue(Value value) const
Finds an LLVM IR value corresponding to the given MLIR value.
SmallVector< llvm::Value * > lookupValues(ValueRange values)
Looks up remapped a list of remapped values.
llvm::Module * getLLVMModule()
Returns the LLVM module in which the IR is being constructed.
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:60
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
InFlightDiagnostic emitError(const Twine &message={})
Emit an error about fatal conditions with this operation, reporting up to any diagnostic handlers tha...
Definition: Operation.cpp:267
OperationName getName()
The name of an operation is the key identifier for it.
Definition: Operation.h:119
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:74
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
Include the generated interface declarations.
Definition: CallGraph.h:229
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:285
void registerOffloadingLLVMTranslationInterfaceExternalModels(mlir::DialectRegistry &registry)
Registers the offloading LLVM translation interfaces for gpu.select_object.
Include the generated interface declarations.
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
Utility class for the GPU dialect to represent triples of Values accessible through ....
Definition: GPUDialect.h:37