Add gpu::GPUDialect::isKernel helper.

Also some mild cleanup of the kernel to cubin conversion pass.

PiperOrigin-RevId: 254959303

1 files changed, 4 insertions, 8 deletions

diff --git a/mlir/lib/Conversion/GPUToCUDA/ConvertKernelFuncToCubin.cpp b/mlir/lib/Conversion/GPUToCUDA/ConvertKernelFuncToCubin.cpp
index 739bdde5330..2ff9435c357 100644
--- a/mlir/lib/Conversion/GPUToCUDA/ConvertKernelFuncToCubin.cpp
+++ b/mlir/lib/Conversion/GPUToCUDA/ConvertKernelFuncToCubin.cpp
@@ -83,7 +83,7 @@ std::string translateModuleToPtx(llvm::Module &module,
 
 using OwnedCubin = std::unique_ptr<std::vector<char>>;
 
-llvm::Optional<OwnedCubin> compilePtxToCubin(std::string ptx,
+llvm::Optional<OwnedCubin> compilePtxToCubin(std::string &ptx,
                                              Function &function) {
   RETURN_ON_CUDA_ERROR(cuInit(0), "cuInit");
 
@@ -161,10 +161,8 @@ LogicalResult translateGpuKernelToCubinAnnotation(Function &function) {
   auto llvmModule = translateModuleToNVVMIR(*module);
   auto cubin = convertModuleToCubin(*llvmModule, function);
 
-  if (!cubin) {
-    function.emitError("Translation to CUDA binary failed.");
-    return failure();
-  }
+  if (!cubin)
+    return function.emitError("Translation to CUDA binary failed.");
 
   function.setAttr(kCubinAnnotation,
                    builder.getStringAttr(
@@ -193,9 +191,7 @@ public:
     LLVMInitializeNVPTXAsmPrinter();
 
     for (auto &function : getModule()) {
-      UnitAttr isKernelAttr = function.getAttrOfType<UnitAttr>(
-          gpu::GPUDialect::getKernelFuncAttrName());
-      if (!isKernelAttr || function.isExternal()) {
+      if (!gpu::GPUDialect::isKernel(&function) || function.isExternal()) {
         continue;
       }
       if (failed(translateGpuKernelToCubinAnnotation(function)))