[flang-commits] [flang] [flang][cuda] Relax the verifier for cuf.register_kernel op (PR #112585)

[Valentin Clement バレンタイン クレメン via flang-commits]

https://github.com/clementval updated https://github.com/llvm/llvm-project/pull/112585

>From b2d4f9b531166699b668afb105a09bb6b8e42cce Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Wed, 16 Oct 2024 10:16:38 -0700
Subject: [PATCH 1/2] [flang][cuda] Relax the verifier for cuf.register_kernel
 op

---
 flang/lib/Optimizer/Dialect/CUF/CUFOps.cpp | 12 ++++++------
 flang/test/Fir/cuf-invalid.fir             | 15 ---------------
 2 files changed, 6 insertions(+), 21 deletions(-)

diff --git a/flang/lib/Optimizer/Dialect/CUF/CUFOps.cpp b/flang/lib/Optimizer/Dialect/CUF/CUFOps.cpp
index 9e3bbd1f9cbee9..5f7232e3dd4b09 100644
--- a/flang/lib/Optimizer/Dialect/CUF/CUFOps.cpp
+++ b/flang/lib/Optimizer/Dialect/CUF/CUFOps.cpp
@@ -281,12 +281,12 @@ mlir::LogicalResult cuf::RegisterKernelOp::verify() {
 
   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");
-
+  if (func) {
+    // Only check if the gpu.func is found. It might be converted to LLVMFuncOp
+    // already.
+    if (!func.isKernel())
+      return emitOpError("only kernel gpu.func can be registered");
+  }
   return mlir::success();
 }
 
diff --git a/flang/test/Fir/cuf-invalid.fir b/flang/test/Fir/cuf-invalid.fir
index a5747b8ee4a3b3..ac488ab3e5abf5 100644
--- a/flang/test/Fir/cuf-invalid.fir
+++ b/flang/test/Fir/cuf-invalid.fir
@@ -143,21 +143,6 @@ module attributes {gpu.container_module} {
 
 // -----
 
-module attributes {gpu.container_module} {
-  gpu.module @cuda_device_mod {
-    gpu.func @_QPsub_device1() {
-      gpu.return
-    }
-  }
-  llvm.func internal @__cudaFortranConstructor() {
-    // expected-error at +1{{'cuf.register_kernel' op device function not found}}
-    cuf.register_kernel @cuda_device_mod::@_QPsub_device2
-    llvm.return
-  }
-}
-
-// -----
-
 module attributes {gpu.container_module} {
   llvm.func internal @__cudaFortranConstructor() {
     // expected-error at +1{{'cuf.register_kernel' op gpu module not found}}

>From 9487980a2088e355c42fbbc8c4f5c92c42eb0bd1 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Wed, 16 Oct 2024 12:19:10 -0700
Subject: [PATCH 2/2] Looks at LLVMFuncOp

---
 flang/lib/Optimizer/Dialect/CUF/CUFOps.cpp | 15 +++++++----
 flang/test/Fir/cuf-invalid.fir             | 30 ++++++++++++++++++++++
 2 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/flang/lib/Optimizer/Dialect/CUF/CUFOps.cpp b/flang/lib/Optimizer/Dialect/CUF/CUFOps.cpp
index 5f7232e3dd4b09..e6fe98791aa30b 100644
--- a/flang/lib/Optimizer/Dialect/CUF/CUFOps.cpp
+++ b/flang/lib/Optimizer/Dialect/CUF/CUFOps.cpp
@@ -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"
@@ -280,14 +281,18 @@ mlir::LogicalResult cuf::RegisterKernelOp::verify() {
     return emitOpError("gpu module not found");
 
   mlir::SymbolTable gpuSymTab(gpuMod);
-  auto func = gpuSymTab.lookup<mlir::gpu::GPUFuncOp>(getKernelName());
-  if (func) {
-    // Only check if the gpu.func is found. It might be converted to LLVMFuncOp
-    // already.
+  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 mlir::success();
+  return emitOpError("device function not found");
 }
 
 // Tablegen operators
diff --git a/flang/test/Fir/cuf-invalid.fir b/flang/test/Fir/cuf-invalid.fir
index ac488ab3e5abf5..8a1eb48576832c 100644
--- a/flang/test/Fir/cuf-invalid.fir
+++ b/flang/test/Fir/cuf-invalid.fir
@@ -143,6 +143,21 @@ module attributes {gpu.container_module} {
 
 // -----
 
+module attributes {gpu.container_module} {
+  gpu.module @cuda_device_mod {
+    gpu.func @_QPsub_device1() {
+      gpu.return
+    }
+  }
+  llvm.func internal @__cudaFortranConstructor() {
+    // expected-error at +1{{'cuf.register_kernel' op device function not found}}
+    cuf.register_kernel @cuda_device_mod::@_QPsub_device2
+    llvm.return
+  }
+}
+
+// -----
+
 module attributes {gpu.container_module} {
   llvm.func internal @__cudaFortranConstructor() {
     // expected-error at +1{{'cuf.register_kernel' op gpu module not found}}
@@ -160,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 at +1{{'cuf.register_kernel' op only gpu.kernel llvm.func can be registered}}
+    cuf.register_kernel @cuda_device_mod::@_QPsub_device1
+    llvm.return
+  }
+}