GPUDialect.h

Go to the documentation of this file.

//===- GPUDialect.h - MLIR Dialect for GPU Kernels --------------*- C++ -*-===//
//
// Part of the MLIR 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 defines the GPU kernel-related operations and puts them in the
// corresponding dialect.
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_DIALECT_GPU_GPUDIALECT_H
#define MLIR_DIALECT_GPU_GPUDIALECT_H

#include "mlir/IR/Dialect.h"
#include "mlir/IR/FunctionSupport.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/OpImplementation.h"
#include "mlir/IR/SymbolTable.h"

namespace mlir {
class FuncOp;

namespace gpu {

class GPUDialect : public Dialect {
public:
  explicit GPUDialect(MLIRContext *context);
  static StringRef getDialectNamespace() { return "gpu"; }

  static StringRef getContainerModuleAttrName() {
    return "gpu.container_module";
  }

  static StringRef getDialectName();

  static StringRef getKernelFuncAttrName() { return "gpu.kernel"; }

  static StringRef getKernelModuleAttrName() { return "gpu.kernel_module"; }

  static bool isKernel(Operation *op);

  // TODO(zinenko,herhut): consider generalizing this.
  static unsigned getNumWorkgroupDimensions() { return 3; }

  static unsigned getWorkgroupAddressSpace() { return 3; }

  static unsigned getPrivateAddressSpace() { return 5; }

  LogicalResult verifyOperationAttribute(Operation *op,
                                         NamedAttribute attr) override;
};

struct KernelDim3 {
  Value x;
  Value y;
  Value z;
};

#define GET_OP_CLASSES
#include "mlir/Dialect/GPU/GPUOps.h.inc"

} // end namespace gpu
} // end namespace mlir

#endif // MLIR_DIALECT_GPU_GPUDIALECT_H