1 //===- GPUDialect.h - MLIR Dialect for GPU Kernels --------------*- C++ -*-===//
2 //
3 // Part of the MLIR 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 defines the GPU kernel-related operations and puts them in the
10 // corresponding dialect.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #ifndef MLIR_DIALECT_GPU_GPUDIALECT_H
15 #define MLIR_DIALECT_GPU_GPUDIALECT_H
16 
17 #include "mlir/IR/Dialect.h"
18 #include "mlir/IR/FunctionSupport.h"
19 #include "mlir/IR/OpDefinition.h"
20 #include "mlir/IR/OpImplementation.h"
21 #include "mlir/IR/SymbolTable.h"
22 
23 namespace mlir {
24 class FuncOp;
25 
26 namespace gpu {
27 
28 /// The dialect containing GPU kernel launching operations and related
29 /// facilities.
30 class GPUDialect : public Dialect {
31 public:
32   /// Create the dialect in the given `context`.
33   explicit GPUDialect(MLIRContext *context);
34   /// Get dialect namespace.
getDialectNamespace()35   static StringRef getDialectNamespace() { return "gpu"; }
36 
37   /// Get the name of the attribute used to annotate the modules that contain
38   /// kernel modules.
getContainerModuleAttrName()39   static StringRef getContainerModuleAttrName() {
40     return "gpu.container_module";
41   }
42 
43   /// Get the canonical string name of the dialect.
44   static StringRef getDialectName();
45 
46   /// Get the name of the attribute used to annotate external kernel functions.
getKernelFuncAttrName()47   static StringRef getKernelFuncAttrName() { return "gpu.kernel"; }
48 
49   /// Get the name of the attribute used to annotate kernel modules.
getKernelModuleAttrName()50   static StringRef getKernelModuleAttrName() { return "gpu.kernel_module"; }
51 
52   /// Returns whether the given function is a kernel function, i.e., has the
53   /// 'gpu.kernel' attribute.
54   static bool isKernel(Operation *op);
55 
56   /// Returns the number of workgroup (thread, block) dimensions supported in
57   /// the GPU dialect.
58   // TODO(zinenko,herhut): consider generalizing this.
getNumWorkgroupDimensions()59   static unsigned getNumWorkgroupDimensions() { return 3; }
60 
61   /// Returns the numeric value used to identify the workgroup memory address
62   /// space.
getWorkgroupAddressSpace()63   static unsigned getWorkgroupAddressSpace() { return 3; }
64 
65   /// Returns the numeric value used to identify the private memory address
66   /// space.
getPrivateAddressSpace()67   static unsigned getPrivateAddressSpace() { return 5; }
68 
69   LogicalResult verifyOperationAttribute(Operation *op,
70                                          NamedAttribute attr) override;
71 };
72 
73 /// Utility class for the GPU dialect to represent triples of `Value`s
74 /// accessible through `.x`, `.y`, and `.z` similarly to CUDA notation.
75 struct KernelDim3 {
76   Value x;
77   Value y;
78   Value z;
79 };
80 
81 #define GET_OP_CLASSES
82 #include "mlir/Dialect/GPU/GPUOps.h.inc"
83 
84 } // end namespace gpu
85 } // end namespace mlir
86 
87 #endif // MLIR_DIALECT_GPU_GPUDIALECT_H
88