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