Skip to content

Commit

Permalink
[flang][cuda] Relax the verifier for cuf.register_kernel op (#112585)
Browse files Browse the repository at this point in the history
Relax the verifier since the `gpu.func` might be converted to
`llvm.func` before `cuf.register_kernel` is converted.
  • Loading branch information
clementval authored Oct 17, 2024
1 parent 98b419c commit 834d001
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 9 deletions.
27 changes: 18 additions & 9 deletions flang/lib/Optimizer/Dialect/CUF/CUFOps.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
#include "flang/Optimizer/Dialect/FIRAttr.h"
#include "flang/Optimizer/Dialect/FIRType.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinOps.h"
Expand Down Expand Up @@ -276,18 +277,26 @@ mlir::LogicalResult cuf::RegisterKernelOp::verify() {

mlir::SymbolTable symTab(mod);
auto gpuMod = symTab.lookup<mlir::gpu::GPUModuleOp>(getKernelModuleName());
if (!gpuMod)
if (!gpuMod) {
// If already a gpu.binary then stop the check here.
if (symTab.lookup<mlir::gpu::BinaryOp>(getKernelModuleName()))
return mlir::success();
return emitOpError("gpu module not found");
}

mlir::SymbolTable gpuSymTab(gpuMod);
auto func = gpuSymTab.lookup<mlir::gpu::GPUFuncOp>(getKernelName());
if (!func)
return emitOpError("device function not found");

if (!func.isKernel())
return emitOpError("only kernel gpu.func can be registered");

return mlir::success();
if (auto func = gpuSymTab.lookup<mlir::gpu::GPUFuncOp>(getKernelName())) {
if (!func.isKernel())
return emitOpError("only kernel gpu.func can be registered");
return mlir::success();
} else if (auto func =
gpuSymTab.lookup<mlir::LLVM::LLVMFuncOp>(getKernelName())) {
if (!func->getAttrOfType<mlir::UnitAttr>(
mlir::gpu::GPUDialect::getKernelFuncAttrName()))
return emitOpError("only gpu.kernel llvm.func can be registered");
return mlir::success();
}
return emitOpError("device function not found");
}

// Tablegen operators
Expand Down
15 changes: 15 additions & 0 deletions flang/test/Fir/cuf-invalid.fir
Original file line number Diff line number Diff line change
Expand Up @@ -175,3 +175,18 @@ module attributes {gpu.container_module} {
llvm.return
}
}

// -----

module attributes {gpu.container_module} {
gpu.module @cuda_device_mod {
llvm.func @_QPsub_device1() {
llvm.return
}
}
llvm.func internal @__cudaFortranConstructor() {
// expected-error@+1{{'cuf.register_kernel' op only gpu.kernel llvm.func can be registered}}
cuf.register_kernel @cuda_device_mod::@_QPsub_device1
llvm.return
}
}

0 comments on commit 834d001

Please sign in to comment.