maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

[mlir] add verifiers for NVVM and ROCDL kernel attributes 

Make sure they can only be attached to LLVM functions as a result of converting
GPU functions to the LLVM Dialect.
This commit is contained in:
Alex Zinenko 2021-02-16 18:05:47 +01:00
parent 627cfd4394
commit 9cd47a26d5
8 changed files with 39 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
View File

@ -24,6 +24,7 @@ def NVVM_Dialect : Dialect {
let name = "nvvm";
let cppNamespace = "::mlir::NVVM";
let dependentDialects = ["LLVM::LLVMDialect"];
let hasOperationAttrVerify = 1;
let extraClassDeclaration = [{
/// Get the name of the attribute used to annotate external kernel

1
mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
View File

@ -24,6 +24,7 @@ def ROCDL_Dialect : Dialect {
let name = "rocdl";
let cppNamespace = "::mlir::ROCDL";
let dependentDialects = ["LLVM::LLVMDialect"];
let hasOperationAttrVerify = 1;
let extraClassDeclaration = [{
/// Get the name of the attribute used to annotate external kernel

12
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
View File

@ -145,5 +145,17 @@ void NVVMDialect::initialize() {
allowUnknownOperations();
}
LogicalResult NVVMDialect::verifyOperationAttribute(Operation *op,
NamedAttribute attr) {
// Kernel function attribute should be attached to functions.
if (attr.first == NVVMDialect::getKernelFuncAttrName()) {
if (!isa<LLVM::LLVMFuncOp>(op)) {
return op->emitError() << "'" << NVVMDialect::getKernelFuncAttrName()
<< "' attribute attached to unexpected op";
}
}
return success();
}
#define GET_OP_CLASSES
#include "mlir/Dialect/LLVMIR/NVVMOps.cpp.inc"

12
mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp
View File

@ -91,5 +91,17 @@ void ROCDLDialect::initialize() {
allowUnknownOperations();
}
LogicalResult ROCDLDialect::verifyOperationAttribute(Operation *op,
NamedAttribute attr) {
// Kernel function attribute should be attached to functions.
if (attr.first == ROCDLDialect::getKernelFuncAttrName()) {
if (!isa<LLVM::LLVMFuncOp>(op)) {
return op->emitError() << "'" << ROCDLDialect::getKernelFuncAttrName()
<< "' attribute attached to unexpected op";
}
}
return success();
}
#define GET_OP_CLASSES
#include "mlir/Dialect/LLVMIR/ROCDLOps.cpp.inc"

5
mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
View File

@ -47,10 +47,7 @@ LogicalResult mlir::NVVMDialectLLVMIRTranslationInterface::amendOperation(
Operation *op, NamedAttribute attribute,
LLVM::ModuleTranslation &moduleTranslation) const {
if (attribute.first == NVVM::NVVMDialect::getKernelFuncAttrName()) {
auto func = dyn_cast<LLVM::LLVMFuncOp>(op);
if (!func)
return failure();
auto func = cast<LLVM::LLVMFuncOp>(op);
llvm::LLVMContext &llvmContext = moduleTranslation.getLLVMContext();
llvm::Function *llvmFunc = moduleTranslation.lookupFunction(func.getName());
llvm::Metadata *llvmMetadata[] = {

4
mlir/lib/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.cpp
View File

@ -54,9 +54,7 @@ LogicalResult mlir::ROCDLDialectLLVMIRTranslationInterface::amendOperation(
Operation *op, NamedAttribute attribute,
LLVM::ModuleTranslation &moduleTranslation) const {
if (attribute.first == ROCDL::ROCDLDialect::getKernelFuncAttrName()) {
auto func = dyn_cast<LLVM::LLVMFuncOp>(op);
if (!func)
return failure();
auto func = cast<LLVM::LLVMFuncOp>(op);
// For GPU kernels,
// 1. Insert AMDGPU_KERNEL calling convention.

7
mlir/test/Dialect/LLVMIR/nvvm.mlir
View File

@ -1,4 +1,4 @@
// RUN: mlir-opt %s | FileCheck %s
// RUN: mlir-opt %s -split-input-file -verify-diagnostics | FileCheck %s
func @nvvm_special_regs() -> i32 {
// CHECK: nvvm.read.ptx.sreg.tid.x : i32
@ -68,3 +68,8 @@ func @nvvm_mma(%a0 : vector<2xf16>, %a1 : vector<2xf16>,
%0 = nvvm.mma.sync %a0, %a1, %b0, %b1, %c0, %c1, %c2, %c3, %c4, %c5, %c6, %c7 {alayout="row", blayout="col"} : (vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, f32, f32, f32, f32, f32, f32, f32, f32) -> !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32)>
llvm.return %0 : !llvm.struct<(f32, f32, f32, f32, f32, f32, f32, f32)>
}
// -----
// expected-error@below {{attribute attached to unexpected op}}
func private @expected_llvm_func() attributes { nvvm.kernel }

6
mlir/test/Dialect/LLVMIR/rocdl.mlir
View File

@ -1,4 +1,4 @@
// RUN: mlir-opt %s | FileCheck %s
// RUN: mlir-opt %s -split-input-file -verify-diagnostics | FileCheck %s
func @rocdl_special_regs() -> i32 {
// CHECK-LABEL: rocdl_special_regs
@ -167,3 +167,7 @@ llvm.func @rocdl.mubuf(%rsrc : vector<4xi32>, %vindex : i32,
llvm.return
}
// -----
// expected-error@below {{attribute attached to unexpected op}}
func private @expected_llvm_func() attributes { rocdl.kernel }