[Mlir-commits] [mlir] [Func][GPU] Create func::ConstantOp using parents with SymbolTable trait (PR #107748)

[Artem Kroviakov]

https://github.com/akroviakov updated https://github.com/llvm/llvm-project/pull/107748

>From bfb43b0289f7e90bf0275c548393b7f16cacd706 Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at tum.de>
Date: Sun, 8 Sep 2024 12:43:34 +0000
Subject: [PATCH] [Func][GPU] Create func::ConstantOp using parents with
 SymbolTable trait

---
 mlir/lib/Dialect/Func/IR/FuncOps.cpp          |  4 +++-
 .../GPU/indirect-device-func-call.mlir        | 21 +++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)
 create mode 100644 mlir/test/Dialect/GPU/indirect-device-func-call.mlir

diff --git a/mlir/lib/Dialect/Func/IR/FuncOps.cpp b/mlir/lib/Dialect/Func/IR/FuncOps.cpp
index c719981769b9e1..f756c64d793fed 100644
--- a/mlir/lib/Dialect/Func/IR/FuncOps.cpp
+++ b/mlir/lib/Dialect/Func/IR/FuncOps.cpp
@@ -128,7 +128,9 @@ LogicalResult ConstantOp::verify() {
   Type type = getType();
 
   // Try to find the referenced function.
-  auto fn = (*this)->getParentOfType<ModuleOp>().lookupSymbol<FuncOp>(fnName);
+  SymbolTable symbolTable(
+      (*this)->getParentWithTrait<mlir::OpTrait::SymbolTable>());
+  auto fn = symbolTable.lookup<FuncOp>(fnName);
   if (!fn)
     return emitOpError() << "reference to undefined function '" << fnName
                          << "'";
diff --git a/mlir/test/Dialect/GPU/indirect-device-func-call.mlir b/mlir/test/Dialect/GPU/indirect-device-func-call.mlir
new file mode 100644
index 00000000000000..91d7f1cd6c67d9
--- /dev/null
+++ b/mlir/test/Dialect/GPU/indirect-device-func-call.mlir
@@ -0,0 +1,21 @@
+// RUN: mlir-opt -test-gpu-rewrite -convert-func-to-llvm %s | FileCheck %s
+
+gpu.module @kernels {
+    // CHECK-LABEL: @hello
+    // CHECK-SAME: %[[ARG0:.*]]: f32
+    func.func @hello(%arg0 : f32) {
+        %tid_x = gpu.thread_id x
+        %csti8 = arith.constant 2 : i8
+        gpu.printf "Hello from %lld, %d, %f\n" %tid_x, %csti8, %arg0  : index, i8, f32
+        return
+    }
+    // CHECK-LABEL: @hello_indirect
+    gpu.func @hello_indirect() kernel { 
+        %cstf32 = arith.constant 3.0 : f32
+        // CHECK: %[[DEVICE_FUNC_ADDR:.*]] = llvm.mlir.addressof @hello : !llvm.ptr
+        %func_ref = func.constant @hello : (f32) -> ()
+        // CHECK: llvm.call %[[DEVICE_FUNC_ADDR]](%{{.*}}) : !llvm.ptr, (f32) -> ()
+        func.call_indirect %func_ref(%cstf32) : (f32) -> ()
+        gpu.return
+    }
+}